Update src/pipeline.py

src/pipeline.py

@@ -46,7 +46,7 @@ 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())
     
     # prune.l1_unstructured(module, 'weight', amount=0.1)
@@ -54,8 +54,7 @@ def load_pipeline() -> Pipeline:
         ckpt_id,
         vae=vae,
         torch_dtype=dtype,
-        ).to(device)
-
+        )
 
     # # List comprehensions to get devices of modules
     # cpu_modules = [name for name, module in pipeline.transformer.named_modules() if next(module.parameters(), None) is None or next(module.parameters()).device == torch.device('cpu')]