Channels last

src/pipeline.py

@@ -32,7 +32,7 @@ def load_pipeline() -> Pipeline:
             path, 
             use_safetensors=False,
             local_files_only=True,
-            torch_dtype=torch.bfloat16).to(memory_format=torch.channels_last)
+            torch_dtype=torch.bfloat16)
     
     pipeline = FluxPipeline.from_pretrained(
         CHECKPOINT,
@@ -41,6 +41,7 @@ def load_pipeline() -> Pipeline:
         local_files_only=True,
         torch_dtype=torch.bfloat16,
     )
+    pipeline.transformer.to(memory_format=torch.channels_last)
     pipeline.to("cuda")
     for _ in range(4):
         pipeline("cat", num_inference_steps=4)