forward to gpu before quant

wfx.py

@@ -86,7 +86,7 @@ class WFX():
         
         self.T2IPipeline = AutoPipelineForText2Image.from_pretrained(args.model, **extra_kwargs)
         self.T2IPipeline.safety_checker = None
-        # self.T2IPipeline.to(torch.device('cuda:0'))
+        self.T2IPipeline.to(torch.device('cuda:0'))
         
         if args.quantize_unet:
             self.T2IPipeline.unet = quantize_unet(self.T2IPipeline.unet)
@@ -94,7 +94,6 @@ class WFX():
         logger.info('compiling model...')
         self.T2IPipeline = compile(self.T2IPipeline, self.compiler_config)
         
-        self.T2IPipeline.to(torch.device('cuda:0'))
         self.warmup()
     
     def warmup(self) -> None:
@@ -103,7 +102,7 @@ class WFX():
         
         warmup_kwargs = dict(
             prompt='a photo of a cat',
-            height=768,
+            height=512,
             width=512,
             num_inference_steps=30,
             generator=torch.Generator(device='cuda:0').manual_seed(0),