From 922db6aa17ec05de9f4d8e6e9799f80384021dc4 Mon Sep 17 00:00:00 2001 From: Pete Date: Thu, 22 Feb 2024 11:54:44 -0800 Subject: [PATCH] Only run the profiler through a single cycle (#463) --- olmo/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/olmo/train.py b/olmo/train.py index f459ad88d..79132f0fc 100644 --- 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"