From 13aaa66ec21c397240a9b972d818430b39112588 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Sat, 2 Aug 2025 12:09:23 -0700
Subject: [PATCH] Make sure context is on the right device. (#9154)

---
 comfy/model_base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index 3ff8106d7..4556ee138 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -171,12 +171,12 @@ class BaseModel(torch.nn.Module):
             dtype = self.manual_cast_dtype
 
         xc = xc.to(dtype)
+        device = xc.device
         t = self.model_sampling.timestep(t).float()
         if context is not None:
-            context = context.to(dtype)
+            context = context.to(dtype=dtype, device=device)
 
         extra_conds = {}
-        device = xc.device
         for o in kwargs:
             extra = kwargs[o]