Fix BumpAllocator error when no input_ids (#5564)

2025-04-20 02:20:53 -07:00
parent 80ac527d22
commit 0a0dd34e6a
2 changed files with 6 additions and 2 deletions
@@ -94,7 +94,9 @@ class DeepseekModelNextN(nn.Module):
        zero_allocator = BumpAllocator(
            buffer_size=2,
            dtype=torch.float32,
-            device=input_ids.device,
+            device=(
+                input_embeds.device if input_embeds is not None else input_ids.device
+            ),
        )

        if input_embeds is None: