Fix BitMarModel class and tensor shapes for main

modeling_bitmar.py

@@ -1537,12 +1537,28 @@ class BitMarModel(PreTrainedModel):
             has_vision: Boolean tensor [batch_size] indicating which samples have real vision features
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # CRITICAL FIX: Ensure input_ids are integers
+        if input_ids.dtype != torch.long:
+            input_ids = input_ids.long()
+        
+        # CRITICAL FIX: Ensure labels are integers if provided
+        if labels is not None and labels.dtype != torch.long:
+            labels = labels.long()
         
         if input_ids is None:
             raise ValueError("input_ids must be provided")
         
         batch_size, seq_len = input_ids.shape
 
+        # Handle missing attention mask
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.float)
+        
+        # Ensure attention_mask is float
+        if attention_mask.dtype != torch.float:
+            attention_mask = attention_mask.float()
+
         # Handle missing vision features
         if vision_features is None:
             vision_features = torch.zeros(batch_size, self.config.vision_encoder_dim,