From f30c554078acad4fb4cf5820e5d30fadf5ba9ed1 Mon Sep 17 00:00:00 2001 From: Akshay Raj Date: Mon, 6 Jul 2026 13:12:54 -0700 Subject: [PATCH] Send T5 tokens to the T5 encoder's own device when caching text embeddings When caching text embeddings for dual text-encoder models (CLIP + T5), the T5 branch sent its input ids to `device`, which is computed from the first encoder (CLIP). If the large T5 encoder is offloaded to a different device (e.g. CPU offload on low-VRAM / unified-memory setups) while CLIP stays on the GPU, the token tensor lands on CLIP's device while T5's weights are elsewhere, raising: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu Send the tokens to the T5 encoder's own device instead. When both encoders are co-located (the common single-GPU case) this is a no-op, so existing setups are unaffected. --- toolkit/train_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/toolkit/train_tools.py b/toolkit/train_tools.py index 78e2183c37..7084f6c198 100644 --- a/toolkit/train_tools.py +++ b/toolkit/train_tools.py @@ -561,7 +561,7 @@ def encode_prompts_flux( ) text_input_ids = text_inputs.input_ids - prompt_embeds = text_encoder[1](text_input_ids.to(device), output_hidden_states=False)[0] + prompt_embeds = text_encoder[1](text_input_ids.to(text_encoder[1].device), output_hidden_states=False)[0] dtype = text_encoder[1].dtype prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)