Hugging Face's logo

Spaces:
AlexTransformer
/
ernie_embebbing_layer_example
Sleeping

App Files Community
AlexTransformer commited on
Commit
c48013e
·
verified ·
1 Parent(s): 029adaf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -11,7 +11,8 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
11
  model = AutoModelForCausalLM.from_pretrained(
12
  model_name,
13
  trust_remote_code=True,
14
- torch_dtype=torch.float32,
 
15
  )
16
  embedding_layer = model.get_input_embeddings()
17
 
 
11
  model = AutoModelForCausalLM.from_pretrained(
12
  model_name,
13
  trust_remote_code=True,
14
+ device_map="auto",
15
+ torch_dtype=torch.bfloat16
16
  )
17
  embedding_layer = model.get_input_embeddings()
18