reformatted (#1129)

huggingface · Dec 23, 2023 · e0dec27 · e0dec27
1 parent 6ef785a
commit e0dec27
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/scripts/ddpo.py b/examples/scripts/ddpo.py
@@ -99,7 +99,7 @@ def __init__(self, *, dtype, model_id, model_filename):
             cached_path = hf_hub_download(model_id, model_filename)
         except EntryNotFoundError:
             cached_path = os.path.join(model_id, model_filename)
-        state_dict = torch.load(cached_path)
+        state_dict = torch.load(cached_path, map_location=torch.device("cpu"))
         self.mlp.load_state_dict(state_dict)
         self.dtype = dtype
         self.eval()