tokenizer

mayank31398 · mayank31398 · commit 78d206619f9e · 2025-06-29T18:46:25.000+05:30
Signed-off-by: Mayank Mishra &lt;mayank31398@gmail.com&gt;
diff --git a/lm_engine/pretrain.py b/lm_engine/pretrain.py
@@ -247,7 +247,9 @@ def train_step_without_pipeline_parallel(
             metrics_tracker = metrics_tracker / gradient_accumulation_steps
 
         metrics_tracker["grad_norm"] = (
-            torch.tensor(0, device=torch.cuda.current_device()) if grad_norm is None else grad_norm
+            torch.zeros((1,), device=torch.cuda.current_device(), dtype=torch.float32)
+            if grad_norm is None
+            else grad_norm
         )
 
         for key in metrics_tracker:

Original file line number	Diff line number	Diff line change
`@@ -247,7 +247,9 @@ def train_step_without_pipeline_parallel(`
`247`	`247`	`metrics_tracker = metrics_tracker / gradient_accumulation_steps`
`248`	`248`
`249`	`249`	`metrics_tracker["grad_norm"] = (`
`250`		`- torch.tensor(0, device=torch.cuda.current_device()) if grad_norm is None else grad_norm`
	`250`	`+ torch.zeros((1,), device=torch.cuda.current_device(), dtype=torch.float32)`
	`251`	`+ if grad_norm is None`
	`252`	`+ else grad_norm`
`251`	`253`	`)`
`252`	`254`
`253`	`255`	`for key in metrics_tracker:`