Update profiler registration check (#6668)

Resolves #5432.
microsoft · Oct 25, 2024 · 54903e0 · 54903e0
1 parent 24285d6
commit 54903e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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):