In rexgen example: Enable training on CUDA-supporting maching and eva…

…luation on non-CUDA-supporting machine. (#131)
awslabs · Jan 19, 2021 · 367c79b · 367c79b
1 parent 89be5a3
commit 367c79b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/reaction_prediction/rexgen_direct/utils.py b/examples/reaction_prediction/rexgen_direct/utils.py
@@ -528,7 +528,7 @@ def prepare_reaction_center(args, reaction_center_config):
             n_layers=reaction_center_config['n_layers'],
             n_tasks=reaction_center_config['n_tasks'])
         reaction_center_model.load_state_dict(
-            torch.load(args['center_model_path'])['model_state_dict'])
+            torch.load(args['center_model_path'], map_location=torch.device('cpu'))['model_state_dict'])
         reaction_center_model = reaction_center_model.to(args['device'])
     reaction_center_model.eval()