sohomghosh commited on
Commit
54ef9da
1 Parent(s): efe7d95

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_read = BERTClass(2)
120
- model_read.load_stat_dict(torch.load('pytorch_model.bin')['model_state_dict'])
 
121
 
122
  tokenizer_read = BertTokenizer.from_pretrained('ProsusAI/finbert')
123
  actual_predictions_read = do_predict(tokenizer_read)
 
117
  return actual_predictions
118
 
119
  model_read = BERTClass(2)
120
+ model_read.to(device)
121
+ model_read.load_stat_dict(torch.load('pytorch_model.bin', map_location=device)['model_state_dict'])
122
 
123
  tokenizer_read = BertTokenizer.from_pretrained('ProsusAI/finbert')
124
  actual_predictions_read = do_predict(tokenizer_read)