From 7f31c4e3acdb6301ffdff23b950229d157cb5f96 Mon Sep 17 00:00:00 2001 From: Logan Adams Date: Fri, 25 Oct 2024 10:09:16 -0700 Subject: [PATCH] Update profiler registration check --- deepspeed/profiling/flops_profiler/profiler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py index 96306184e42c..f87f1beb7e4e 100644 --- a/deepspeed/profiling/flops_profiler/profiler.py +++ b/deepspeed/profiling/flops_profiler/profiler.py @@ -115,7 +115,7 @@ def start_time_hook(module, input): get_accelerator().synchronize() module.__start_time__ = time.time() - if not hasattr(module, "__start_time_hook_handle"): + if not hasattr(module, "__start_time_hook_handle__"): module.__start_time_hook_handle__ = module.register_forward_pre_hook(start_time_hook) def end_time_hook(module, input, output):