clean up

allenai · Apr 30, 2024 · 6f136b6 · 6f136b6
1 parent a9f8194
commit 6f136b6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/benchmarks/fsdp/train.py b/src/benchmarks/fsdp/train.py
@@ -117,8 +117,8 @@ def on_trace_ready(p):
             print_rank0(
                 f"Batch [{i+1}/{num_batches}]:\n"
                 f"  loss={loss.item():.3f}\n"
-                f"  throughput/seconds_per_batch={batch_time:.3f}\n",
-                f"  grad/total_norm={norm_str}",
+                f"  throughput/seconds_per_batch={batch_time:.3f}\n"
+                f"  grad/total_norm={norm_str}"
             )
 
             if profile and i == 2: