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

2026-02-13 21:44:31 +01:00 · 2022-07-07 10:08:33 -07:00
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

--- a/dalle2_pytorch/version.py
+++ b/dalle2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.16.16'
+__version__ = '0.16.17'