feat: updated .to() override to handle kwargs

modeling_bert.py

@@ -422,9 +422,9 @@ class BertModel(BertPreTrainedModel):
             pooler_output=pooled_output,
         )
 
-    def to(self, target):
-        result = super().to(target)
-        if isinstance(target, torch.dtype):
+    def to(self, *args, **kwargs):
+        result = super().to(*args, **kwargs)
+        if (len(args) > 0 and isinstance(args[0], torch.dtype)) or "dtype" in kwargs:
             for layer in result.encoder.layers:
                 layer.mixer.inner_cross_attn.alibi_slopes = layer.mixer.inner_cross_attn.alibi_slopes.to(torch.float32)
                 layer.mixer.inner_attn.alibi_slopes = layer.mixer.inner_attn.alibi_slopes.to(torch.float32)