fix script (#741)

mosaicml · Nov 16, 2023 · e796218 · e796218
1 parent db279d0
commit e796218
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/scripts/inference/hf_generate.py b/scripts/inference/hf_generate.py
@@ -217,6 +217,7 @@ def main(args: Namespace) -> None:
         if device is not None:
             print(f'Placing model on {device=}...')
             model.to(device)
+        model.to(model_dtype)
     except Exception as e:
         raise RuntimeError(
             'Unable to load HF model. ' +