Add cloudpickle wrapper for compatibility

rlworkgroup · May 4, 2021 · a9e083a · a9e083a
1 parent 3d4368d
commit a9e083a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -13,7 +13,7 @@
     'akro>=0.0.8',
     'click>=2.0',
     # Older versions don't work with torch.save
-    'cloudpickle>=1.6.0',
+    'cloudpickle',
     'cma==2.7.0',
     'dowel==0.0.3',
     'numpy>=1.14.5',

diff --git a/src/garage/experiment/snapshotter.py b/src/garage/experiment/snapshotter.py
@@ -136,8 +136,12 @@ def save_itr_params(self, itr, params):
 
         if file_name:
             if torch:
+
+                class _pickle_module:
+                    Pickler = cloudpickle.CloudPickler
+
                 params['global_device'] = global_device()
-                torch.save(params, file_name, pickle_module=cloudpickle)
+                torch.save(params, file_name, pickle_module=_pickle_module)
             else:
                 with open(file_name, 'wb') as file:
                     cloudpickle.dump(params, file)