temporarily hardcode torch_dtype=float16 for SimpleGenerativeModel

ArneBinder · Feb 6, 2025 · 520d6ab · 520d6ab
1 parent 2e4fd26
commit 520d6ab
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/pie_modules/models/simple_generative.py b/src/pie_modules/models/simple_generative.py
@@ -81,6 +81,7 @@ def __init__(
         # Note: We do not set expected_super_type=PreTrainedModel for resolve_type() because
         #   AutoModel* classed such as AutoModelForSeq2SeqLM do not inherit from that.
         resolved_base_model_type: Type[PreTrainedModel] = resolve_type(base_model_type)
+        base_model_config["torch_dtype"] = torch.float16
         self.model = resolved_base_model_type.from_pretrained(**base_model_config)
         self.generation_config = self.configure_generation(**(override_generation_kwargs or {}))