From 7c0cd05c272ecb5067e85f0a667fbbcce6b36727 Mon Sep 17 00:00:00 2001 From: Ryan Burgert Date: Sun, 4 Sep 2022 22:11:35 -0400 Subject: [PATCH] Trainer.load can use GPU's other than cuda:0 --- denoising_diffusion_pytorch/denoising_diffusion_pytorch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 2e73fa0..c8d0086 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -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'])