robinhad commited on
Commit
be676b9
1 Parent(s): 13e12b6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -2,14 +2,16 @@ from peft import PeftModel
2
  from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
3
  import gradio as gr
4
  from torch.cuda import is_available
 
5
 
6
- if True:#is_available():
7
  options = dict(
8
  load_in_8bit=True,
9
  device_map="auto",
10
  )
11
  else:
12
  options = {
 
13
  #"low_cpu_mem_usage": True,
14
  #"offload_state_dict": True,
15
  #"offload_folder": "offload",
 
2
  from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
3
  import gradio as gr
4
  from torch.cuda import is_available
5
+ from torch import bfloat16
6
 
7
+ if is_available():
8
  options = dict(
9
  load_in_8bit=True,
10
  device_map="auto",
11
  )
12
  else:
13
  options = {
14
+ "dtype": bfloat16
15
  #"low_cpu_mem_usage": True,
16
  #"offload_state_dict": True,
17
  #"offload_folder": "offload",