Update README.md

README.md

@@ -82,12 +82,16 @@ if __name__ == "__main__":
     # get datasets
     raw_datasets = load_dataset("mila-intel/ProtST-SubcellularLocalization", cache_dir="~/.cache/huggingface/datasets", split='test') # cache_dir defaults to "~/.cache/huggingface/datasets"
 
-    device = torch.device("cpu")
+-   device = torch.device("cpu")
++   device = torch.device("hpu")
     
     protst_model = AutoModel.from_pretrained("mila-intel/ProtST-esm1b", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
     protein_model = protst_model.protein_model
     text_model = protst_model.text_model
     logit_scale = protst_model.logit_scale
++   from habana_frameworks.torch.hpu import wrap_in_hpu_graph
++   protein_model = wrap_in_hpu_graph(protein_model)
++   text_model = wrap_in_hpu_graph(text_model)
     logit_scale.requires_grad = False
     logit_scale = logit_scale.to(device)
     logit_scale = logit_scale.exp()