fix: dtype

modeling_hf_nomic_bert.py

@@ -391,6 +391,7 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         num_labels = kwargs.pop("num_labels", None)
         rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
         strict = kwargs.pop("strict", True)
+        dtype = kwargs.pop("dtype", None)
         if rotary_scaling_factor:
             config.rotary_scaling_factor = rotary_scaling_factor
 
@@ -406,6 +407,9 @@ class NomicBertPreTrainedModel(PreTrainedModel):
                 model = cls(config, *inputs, add_pooling_layer=False)
             else:
                 model = cls(config, *inputs)
+
+        if dtype is not None:
+            model = model.to(dtype=dtype)
         # TODO: fix this
         # Assuming we know what we're doing when loading from disk
         # Prob a bad assumption but i'm tired and want to train this asap
@@ -424,7 +428,7 @@ class NomicBertPreTrainedModel(PreTrainedModel):
             load_return = model.load_state_dict(state_dict, strict=False)
         else:
             # TODO: can probably check config class and see if we need to remap from a bert model
-            state_dict = state_dict_from_pretrained(model_name)
+            state_dict = state_dict_from_pretrained(model_name, dtype=dtype)
             state_dict = remap_bert_state_dict(
                 state_dict,
                 config,