kuleshov-group
/

PlantCaduceus_l24

Feature Extraction

Model card Files Files and versions Community

Jingjing Zhai commited on Jun 10

Commit

208dd61

•

1 Parent(s): c0420f1

Update README

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -17,7 +17,7 @@ from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer
 import torch
 model_path = 'kuleshov-group/PlantCaduceus_l24'
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
-model = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True).to(device)
 model.eval()
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

 import torch
 model_path = 'kuleshov-group/PlantCaduceus_l24'
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True, device_map=device)
 model.eval()
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)