Update modelling_cxrrg.py

modelling_cxrrg.py

@@ -541,4 +541,4 @@ class CXRRGModel(VisionEncoderDecoderModel):
         causal_2d_attention_mask = causal_2d_attention_mask[:, None, None, :]
 
         mixed_causality_4d_attention_mask = torch.cat((non_causal_2d_attention_mask, causal_2d_attention_mask), dim=-1)
-        return mixed_causality_4d_attention_mask
+        return mixed_causality_4d_attention_mask.to(dtype=torch.float)