Only run the profiler through a single cycle (#463)

allenai · Feb 22, 2024 · 922db6a · 922db6a
1 parent 37ca789
commit 922db6a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/train.py b/olmo/train.py
@@ -877,7 +877,7 @@ def fit(self):
         if self.cfg.torch_profiling and get_global_rank() == 0:
             from torch.profiler import schedule
 
-            profiling_schedule = schedule(wait=1, warmup=5, active=3)
+            profiling_schedule = schedule(wait=1, warmup=5, active=3, repeat=1)
 
             def on_trace_ready(p):
                 profiler_output_dir = Path(self.cfg.save_folder) / "profiler"