Hugging Face's logo

kuleshov-group
/
PlantCaduceus_l20

Feature Extraction
Transformers
PyTorch
caduceus
custom_code
Model card Files Community
3
Jingjing Zhai commited on
Commit
4a3df3e
·
1 Parent(s): 3613a7b

Update README

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -17,7 +17,7 @@ from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer
17
  import torch
18
  model_path = 'kuleshov-group/PlantCaduceus_l20'
19
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
20
- model = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True).to(device)
21
  model.eval()
22
  tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
23
 
 
17
  import torch
18
  model_path = 'kuleshov-group/PlantCaduceus_l20'
19
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
20
+ model = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True, device_map=device)
21
  model.eval()
22
  tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
23