Update pipeline.py

pipeline.py

@@ -241,6 +241,7 @@ class FluxWithCFGPipeline(StableDiffusion3Pipeline):
         return self.image_processor.postprocess(image, output_type=output_type)[0]
         
 class FluxWithCFGPipeline(StableDiffusion3Pipeline):
+    @torch.inference_mode()
     def __init__(
         self,
         transformer: FluxTransformer2DModel,
@@ -273,7 +274,7 @@ class FluxWithCFGPipeline(StableDiffusion3Pipeline):
             self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
         )
         self.default_sample_size = 64
-@torch.inference_mode()
+    
     def generate_image(
         self,
         prompt: Union[str, List[str]] = None,