Fix device placement when loading

ikostrikov · Feb 4, 2021 · b25c006 · b25c006
1 parent b1152f8
commit b25c006
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/enjoy.py b/enjoy.py
@@ -50,7 +50,8 @@
 
 # We need to use the same statistics for normalization as used in training
 actor_critic, ob_rms = \
-            torch.load(os.path.join(args.load_dir, args.env_name + ".pt"))
+            torch.load(os.path.join(args.load_dir, args.env_name + ".pt"),
+                        map_location='cpu')
 
 vec_norm = get_vec_normalize(env)
 if vec_norm is not None: