correct modeling_ddpm

addc43af · Patrick von Platen · f9a4532f · addc43af
Commit addc43af authored Jun 07, 2022 by Patrick von Platen
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

models/vision/ddpm/modeling_ddpm.py models/vision/ddpm/modeling_ddpm.py +3 -2

No files found.
--- a/models/vision/ddpm/modeling_ddpm.py
+++ b/models/vision/ddpm/modeling_ddpm.py
@@ -27,7 +27,8 @@ class DDPM(DiffusionPipeline):
        super().__init__()
        self.register_modules(unet=unet, noise_scheduler=noise_scheduler)

-    def __call__(self, generator=None, torch_device=None):
+    def __call__(self, batch_size=1, generator=None, torch_device=None):
+        if torch_device is None:
            torch_device = "cuda" if torch.cuda.is_available() else "cpu"

        self.unet.to(torch_device)