diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py
index 7b3717119e..147ffadfac 100644
--- a/examples/mlperf/model_train.py
+++ b/examples/mlperf/model_train.py
@@ -1397,7 +1397,7 @@ def train_llama3():
 
   if getenv("FAKEDATA"):
     for v in get_parameters(model):
-      v = v.assign(Tensor.empty(v.shape))
+      v = v.assign(Tensor.empty(v.shape, dtype=v.dtype))
 
   is_dp = (DP := getenv("DP", 1)) > 1
   is_mp = (MP := getenv("MP", 1)) > 1