0.27.9

Merge pull request #85 from RyannDaGreat/main
Trainer.load can use GPU's other than cuda:0
2026-09-10 12:01:08 +08:00 · 2022-09-05 02:43:11 -07:00 · 2022-09-05 02:42:49 -07:00 · 2022-09-04 22:11:35 -04:00
2 changed files with 5 additions and 2 deletions
@@ -810,7 +810,10 @@ class Trainer(object):
        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))

    def load(self, milestone):
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'))
+        accelerator = self.accelerator
+        device = accelerator.device
+
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)

        model = self.accelerator.unwrap_model(self.model)
        model.load_state_dict(data['model'])
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
  name = 'denoising-diffusion-pytorch',
  packages = find_packages(),
-  version = '0.27.8',
+  version = '0.27.9',
  license='MIT',
  description = 'Denoising Diffusion Probabilistic Models - Pytorch',
  author = 'Phil Wang',
Author	SHA1	Message	Date
Phil Wang	e0a1bed31a	0.27.9	2022-09-05 02:43:11 -07:00
Phil Wang	82b67fc00a	Merge pull request #85 from RyannDaGreat/main Trainer.load can use GPU's other than cuda:0	2022-09-05 02:42:49 -07:00
Ryan Burgert	7c0cd05c27	Trainer.load can use GPU's other than cuda:0	2022-09-04 22:11:35 -04:00