zachlopez commited on
Commit
5b1c0e6
1 Parent(s): 795dec4

moved to device

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -800,7 +800,7 @@ model = GPT2LMHeadModel.from_pretrained(
800
  output_hidden_states=True
801
  )
802
  if model_fp != None and model_fp != "":
803
- model.load_state_dict(torch.load(model_fp))
804
  model.to(device)
805
  model.eval()
806
 
 
800
  output_hidden_states=True
801
  )
802
  if model_fp != None and model_fp != "":
803
+ model.load_state_dict(torch.load(model_fp, map_location=device))
804
  model.to(device)
805
  model.eval()
806