Fix: explicitly cast to bool tensor before masked scatter
Browse files- modeling_deepseekocr2.py +4 -1
modeling_deepseekocr2.py
CHANGED
|
@@ -497,7 +497,10 @@ class DeepseekOCR2Model(DeepseekV2Model):
|
|
| 497 |
device = torch.device("cuda")
|
| 498 |
else:
|
| 499 |
device = torch.device("cpu")
|
| 500 |
-
|
|
|
|
|
|
|
|
|
|
| 501 |
|
| 502 |
idx += 1
|
| 503 |
|
|
|
|
| 497 |
device = torch.device("cuda")
|
| 498 |
else:
|
| 499 |
device = torch.device("cpu")
|
| 500 |
+
|
| 501 |
+
mask = images_seq_mask[idx].unsqueeze(-1).to(device=device, dtype=torch.bool)
|
| 502 |
+
|
| 503 |
+
inputs_embeds[idx].masked_scatter_(mask, images_in_this_batch)
|
| 504 |
|
| 505 |
idx += 1
|
| 506 |
|