Update logic.py
Browse files
logic.py
CHANGED
@@ -434,7 +434,7 @@ def load_model(train_data, val_data, test_data):
|
|
434 |
model = Model(hidden_channels=32)
|
435 |
with torch.no_grad():
|
436 |
model.encoder(train_data.x_dict, train_data.edge_index_dict)
|
437 |
-
model.load_state_dict(torch.load('model.pt'))
|
438 |
model.eval()
|
439 |
return model
|
440 |
|
|
|
434 |
model = Model(hidden_channels=32)
|
435 |
with torch.no_grad():
|
436 |
model.encoder(train_data.x_dict, train_data.edge_index_dict)
|
437 |
+
model.load_state_dict(torch.load('model.pt',map_location=torch.device('cpu')))
|
438 |
model.eval()
|
439 |
return model
|
440 |
|