Fix CPU/CUDA device mismatch in Klein edit control image encoding (#742)
When training Klein models with a `control_path` (edit/kontext-style
paired datasets), `encode_image_refs()` returns tensors that reside on
the VAE's device (CPU, since the VAE weights are loaded via
`load_file(..., device="cpu")` and are never explicitly moved to the
training device). Concatenating those CPU tensors with the training
latents (`packed_latents`) that live on CUDA raises:
RuntimeError: Expected all tensors to be on the same device
Fix: move `img_cond_seq` and `img_cond_seq_ids` to the same device
(and dtype) as `img_input` / `img_input_ids` before concatenation.
Co-authored-by: HuangYuChuh <HuangYuChuh@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
89d2090962
commit
489b194231
|
|
@ -412,8 +412,8 @@ class Flux2Model(BaseModel):
|
|||
assert img_cond_seq_ids is not None, (
|
||||
"You need to provide either both or neither of the sequence conditioning"
|
||||
)
|
||||
img_input = torch.cat((img_input, img_cond_seq), dim=1)
|
||||
img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids), dim=1)
|
||||
img_input = torch.cat((img_input, img_cond_seq.to(img_input.device, img_input.dtype)), dim=1)
|
||||
img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids.to(img_input_ids.device)), dim=1)
|
||||
|
||||
guidance_vec = torch.full(
|
||||
(img_input.shape[0],),
|
||||
|
|
|
|||
Loading…
Reference in New Issue