Update: cast to device instead of explicitly cude

modeling_deepseekocr2.py

@@ -491,8 +491,13 @@ class DeepseekOCR2Model(DeepseekV2Model):
                 if images_in_this_batch:
                     images_in_this_batch = torch.cat(images_in_this_batch, dim=0)
                     # exit()
-
-                    inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)
+                    if torch.backends.mps.is_available():
+                        device = torch.device("mps")
+                    elif torch.cuda.is_available():
+                        device = torch.device("cuda")
+                    else:
+                        device = torch.device("cpu")
+                    inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).to(device), images_in_this_batch)
 
                 idx += 1