diff --git a/Code/python/notebooks/predict.py b/Code/python/notebooks/predict.py index df8f885..761e331 100644 --- a/Code/python/notebooks/predict.py +++ b/Code/python/notebooks/predict.py @@ -50,7 +50,7 @@ model = cls( device=DEVICE, ) -model.load_state_dict(torch.load("LSTM_wu900_ps150_aw0.001.pt")) +model.load_state_dict(torch.load("LSTM_wu900_ps150_aw0.001.pt", map_location=torch.device('cpu'))) model.to(DEVICE) model.eval() with torch.no_grad():