reformatted (huggingface#1129)

lapp0 · May 10, 2024 · 55cb862 · 55cb862
1 parent 1b9b397
commit 55cb862
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()