Update src/pipeline.py

src/pipeline.py

@@ -29,13 +29,13 @@ def load_pipeline() -> Pipeline:
     dtype, device = torch.bfloat16, "cuda"
     vae = AutoencoderKL.from_pretrained(
         ckpt_id, subfolder="vae", torch_dtype=torch.bfloat16
-    )
+    ).to(device)
     # quantize_(vae, int8_weight_only())
     pipeline = DiffusionPipeline.from_pretrained(
         ckpt_id,
         vae=vae,
         torch_dtype=dtype,
-        )
+        ).to(device)
     for name, module in pipeline.transformer.named_modules():
       if isinstance(module, torch.nn.Linear):
           prune.l1_unstructured(module, 'weight', amount=0.2)