From 403f9e686046f121f6cfaaeac646116c0f5ca275 Mon Sep 17 00:00:00 2001 From: lucidrains Date: Mon, 7 Oct 2024 10:11:15 -0700 Subject: [PATCH] address a warning --- denoising_diffusion_pytorch/denoising_diffusion_pytorch.py | 2 +- denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py | 2 +- denoising_diffusion_pytorch/guided_diffusion.py | 2 +- denoising_diffusion_pytorch/repaint.py | 2 +- denoising_diffusion_pytorch/version.py | 2 +- setup.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 9b9d889cc..4d4f98fc3 100644 --- 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 index 8f5c55ee2..d49d89519 100644 --- 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 index 6e598885a..e610ae4f8 100644 --- 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 index 147564842..dbc66455a 100644 --- 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 index 2b6bf429f..a33997dd1 100644 --- 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 index 1118bb080..333939a73 100644 --- a/setup.py +++ b/setup.py @@ -24,7 +24,7 @@ 'pillow', 'pytorch-fid', 'scipy', - 'torch', + 'torch>=2.0', 'torchvision', 'tqdm' ],