make fp32 path in .py code work correctly

karpathy · Sep 17, 2024 · 685617f · 685617f
1 parent 508c474
commit 685617f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/train_llama3.py b/train_llama3.py
@@ -1101,6 +1101,10 @@ def print0(*args, **kwargs):
         assert args.tokenizer_path is not None and os.path.exists(args.tokenizer_path), f"llama3 tokenizer path {args.tokenizer_path} does not exist"
         model = LLaMA.from_pretrained_llama3_meta(args.ckpt_dir, args.tokenizer_path)
 
+    # convert the model to the desired precision
+    if args.dtype == "float32":
+        model = model.to(torch.float32)
+
     model.train()
     if args.compile:
         if hasattr(config, "coordinate_descent_tuning"):