diffusers
diffusers copied to clipboard
move untruncated_ids to GPU to fix tensor device mismatch error
in SDXL-LoRA-RNPD.ipynb, when External_Captions is set to TRUE and captions exceed 77 tokens, it will cause a device mismatch and result in not training the TE for all the epochs that have been specified (it stops after 1 epoch), and instead move on to train the Unet. the text_input_ids tensor is moved to GPU (as seen in the line text_input_ids = text_inputs.input_ids.to(text_encoder.device)), while untruncated_ids stays on the CPU.
to fix the error untruncated_ids must be moved to the GPU, the same device as text_input_ids, before the comparison:
untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
untruncated_ids = untruncated_ids.to('cuda') # Move untruncated_ids to GPU
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):