address a warning

lucidrains · Oct 7, 2024 · 403f9e6 · 403f9e6
1 parent 7c1a4cf
commit 403f9e6
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 6 deletions.
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -1024,7 +1024,7 @@ def load(self, milestone):
         accelerator = self.accelerator
         device = accelerator.device
 
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device, weights_only=True)
 
         model = self.accelerator.unwrap_model(self.model)
         model.load_state_dict(data['model'])

diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
@@ -815,7 +815,7 @@ def load(self, milestone):
         accelerator = self.accelerator
         device = accelerator.device
 
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device, weights_only=True)
 
         model = self.accelerator.unwrap_model(self.model)
         model.load_state_dict(data['model'])

diff --git a/denoising_diffusion_pytorch/guided_diffusion.py b/denoising_diffusion_pytorch/guided_diffusion.py
@@ -890,7 +890,7 @@ def load(self, milestone):
         accelerator = self.accelerator
         device = accelerator.device
 
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device, weights_only=True)
 
         model = self.accelerator.unwrap_model(self.model)
         model.load_state_dict(data['model'])

diff --git a/denoising_diffusion_pytorch/repaint.py b/denoising_diffusion_pytorch/repaint.py
@@ -1060,7 +1060,7 @@ def load(self, milestone):
         accelerator = self.accelerator
         device = accelerator.device
 
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device, weights_only=True)
 
         model = self.accelerator.unwrap_model(self.model)
         model.load_state_dict(data['model'])

diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py
@@ -1 +1 @@
-__version__ = '2.0.18'
+__version__ = '2.1.0'
diff --git a/setup.py b/setup.py
@@ -24,7 +24,7 @@
     'pillow',
     'pytorch-fid',
     'scipy',
-    'torch',
+    'torch>=2.0',
     'torchvision',
     'tqdm'
   ],