s-conia commited on
Commit
4db788d
1 Parent(s): 75f2b84

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -69,7 +69,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
69
  tokenizer = tr.AutoTokenizer.from_pretrained("sapienzanlp/modello-italia-9b")
70
  model = tr.AutoModelForCausalLM.from_pretrained(
71
  "sapienzanlp/modello-italia-9b",
72
- device_map="auto",
73
  torch_dtype=torch.bfloat16
74
  )
75
 
 
69
  tokenizer = tr.AutoTokenizer.from_pretrained("sapienzanlp/modello-italia-9b")
70
  model = tr.AutoModelForCausalLM.from_pretrained(
71
  "sapienzanlp/modello-italia-9b",
72
+ device_map=device,
73
  torch_dtype=torch.bfloat16
74
  )
75