extra insurance that diffusion prior is on the correct device, when using trainer with accelerator or device was given

2025-12-19 17:54:20 +01:00 · 2022-07-07 10:08:33 -07:00
parent 46be8c32d3
commit ec7cab01d9
2 changed files with 2 additions and 1 deletions
--- a/dalle2_pytorch/trainer.py
+++ b/dalle2_pytorch/trainer.py
@@ -192,6 +192,7 @@ class DiffusionPriorTrainer(nn.Module):
            self.device = diffusion_prior_device
        else:
            self.device = accelerator.device if exists(accelerator) else device
+            diffusion_prior.to(self.device)

        # save model