Update nodes.py

2024-08-28 00:58:20 +03:00
parent 58fee8e567
commit cb5465ad20
1 changed files with 1 additions and 1 deletions
@@ -76,7 +76,7 @@ class DownloadAndLoadCogVideoModel:
            torch._inductor.config.coordinate_descent_tuning = True
            torch._inductor.config.epilogue_fusion = False
            torch._inductor.config.coordinate_descent_check_all_directions = True
-            pipe.transformer.to(memory_format=torch.channels_last)
+            pipe.transformer.to(device).to(memory_format=torch.channels_last)
            pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)

        pipeline = {