use map location to cpu when torch.load

cmu-mlsp · Feb 17, 2023 · aaf0d2d · aaf0d2d
1 parent 0abb04d
commit aaf0d2d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/audiolm_pytorch/trainer.py b/audiolm_pytorch/trainer.py
@@ -262,7 +262,7 @@ def unwrapped_soundstream(self):
     def load(self, path):
         path = Path(path)
         assert path.exists()
-        pkg = torch.load(str(path))
+        pkg = torch.load(str(path), map_location = 'cpu')
 
         self.unwrapped_soundstream.load_state_dict(pkg['model'])
 
@@ -585,7 +585,7 @@ def save(self, path):
     def load(self, path):
         path = Path(path)
         assert path.exists()
-        pkg = torch.load(str(path))
+        pkg = torch.load(str(path), map_location = 'cpu')
 
         transformer = self.accelerator.unwrap_model(self.transformer)
         transformer.load_state_dict(pkg['model'])
@@ -822,7 +822,7 @@ def save(self, path):
     def load(self, path):
         path = Path(path)
         assert path.exists()
-        pkg = torch.load(str(path))
+        pkg = torch.load(str(path), map_location = 'cpu')
 
         transformer = self.accelerator.unwrap_model(self.transformer)
         transformer.load_state_dict(pkg['model'])
@@ -1054,7 +1054,7 @@ def save(self, path):
     def load(self, path):
         path = Path(path)
         assert path.exists()
-        pkg = torch.load(str(path))
+        pkg = torch.load(str(path), map_location = 'cpu')
 
         transformer = self.accelerator.unwrap_model(self.transformer)
         transformer.load_state_dict(pkg['model'])

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'audiolm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.14.0',
+  version = '0.14.1',
   license='MIT',
   description = 'AudioLM - Language Modeling Approach to Audio Generation from Google Research - Pytorch',
   author = 'Phil Wang',