From cb5465ad20fedc0e2d2cf9870424d2519df990e4 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 28 Aug 2024 00:58:20 +0300
Subject: [PATCH] Update nodes.py

---
 nodes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 62f827b..403fd5f 100644
--- a/nodes.py
+++ b/nodes.py
@@ -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 = {