From e451fbbac4268c5e295459b090ba4559b8cd080f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Mon, 13 Jul 2026 14:39:17 +0900 Subject: [PATCH] Create PixArt dropout masks on the input device --- comfy/ldm/pixart/blocks.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/pixart/blocks.py b/comfy/ldm/pixart/blocks.py index 2225076e5..58d838613 100644 --- a/comfy/ldm/pixart/blocks.py +++ b/comfy/ldm/pixart/blocks.py @@ -296,7 +296,7 @@ class LabelEmbedder(nn.Module): Drops labels to enable classifier-free guidance. """ if force_drop_ids is None: - drop_ids = torch.rand(labels.shape[0]).cuda() < self.dropout_prob + drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob else: drop_ids = force_drop_ids == 1 labels = torch.where(drop_ids, self.num_classes, labels) @@ -328,7 +328,7 @@ class CaptionEmbedder(nn.Module): Drops labels to enable classifier-free guidance. """ if force_drop_ids is None: - drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob + drop_ids = torch.rand(caption.shape[0], device=caption.device) < self.uncond_prob else: drop_ids = force_drop_ids == 1 caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption) @@ -363,7 +363,7 @@ class CaptionEmbedderDoubleBr(nn.Module): Drops labels to enable classifier-free guidance. """ if force_drop_ids is None: - drop_ids = torch.rand(global_caption.shape[0]).cuda() < self.uncond_prob + drop_ids = torch.rand(global_caption.shape[0], device=global_caption.device) < self.uncond_prob else: drop_ids = force_drop_ids == 1 global_caption = torch.where(drop_ids[:, None], self.embedding, global_caption)