Use set_default_dtype for uniform dtype init, assert batch size match in CrossAttention

modeling_slip.py

@@ -513,6 +513,11 @@ class SLIPModel(SLIPPreTrainedModel):
     def __init__(self, config: SLIPConfig):
         super().__init__(config)
 
+        # Use config.torch_dtype (set by HF from_pretrained or config.json)
+        dtype = getattr(config, "torch_dtype", None) or torch.float32
+        prev_dtype = torch.get_default_dtype()
+        torch.set_default_dtype(dtype)
+
         # Sensor encoder
         sensor_cfg = config.sensor_encoder
         self.sensor_encoder = SensorTransformerModel(**sensor_cfg)
@@ -544,6 +549,8 @@ class SLIPModel(SLIPPreTrainedModel):
         self.temperature = nn.Parameter(torch.tensor(math.log(1 / 0.07)))
         self.temperature_max = math.log(1 / 0.07)
 
+        torch.set_default_dtype(prev_dtype)
+
     def embed_sensor(self, sensors, sensor_attn_mask=None, time_index=None):
         sensor_tokens, attn_mask = self.sensor_encoder(sensors, sensor_attn_mask, time_index=time_index)
         if hasattr(self, "img_attn_pool"):