TillCyrill
commited on
Commit
•
6646de2
1
Parent(s):
1372a90
device
Browse files
main.py
CHANGED
@@ -42,10 +42,10 @@ atom_mapping = {0:'H', 1:'C', 2:'N', 3:'O', 4:'F', 5:'P', 6:'S', 7:'CL', 8:'BR',
|
|
42 |
model = GNN_MD(11, 64)
|
43 |
state_dict = torch.load(
|
44 |
"best_weights_rep0.pt",
|
45 |
-
map_location=torch.device("
|
46 |
)["model_state_dict"]
|
47 |
model.load_state_dict(state_dict)
|
48 |
-
model = model.to(
|
49 |
model.eval()
|
50 |
|
51 |
|
|
|
42 |
model = GNN_MD(11, 64)
|
43 |
state_dict = torch.load(
|
44 |
"best_weights_rep0.pt",
|
45 |
+
map_location=torch.device("cpu"),
|
46 |
)["model_state_dict"]
|
47 |
model.load_state_dict(state_dict)
|
48 |
+
model = model.to('cpu')
|
49 |
model.eval()
|
50 |
|
51 |
|