kuleshov-group
/

PlantCaduceus_l32

Feature Extraction

Model card Files Files and versions Community

Jingjing Zhai commited on Jun 10

Commit

3d82951

•

1 Parent(s): 35fa607

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_l32'
 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_l32'
 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)