sohomghosh commited on
Commit
8a57cfb
1 Parent(s): 7ffaed7

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -1
README.md CHANGED
@@ -117,7 +117,8 @@ def do_predict(tokenizer):
117
  return actual_predictions
118
 
119
  model_sus = BERTClass(2)
120
- model_sus.load_stat_dict(torch.load('pytorch_model.bin')['model_state_dict'])
 
121
 
122
  tokenizer_sus = BertTokenizer.from_pretrained('roberta-base')
123
  actual_predictions_sus = do_predict(tokenizer_sus)
 
117
  return actual_predictions
118
 
119
  model_sus = BERTClass(2)
120
+ model_sustain.to(device)
121
+ model_sustain.load_state_dict(torch.load('pytorch_model.bin', map_location=device)['model_state_dict'])
122
 
123
  tokenizer_sus = BertTokenizer.from_pretrained('roberta-base')
124
  actual_predictions_sus = do_predict(tokenizer_sus)