Update flops profiler to recurse (#4374)

Fixes #4334.
microsoft · Nov 21, 2023 · 6139122 · 6139122
1 parent a603a21
commit 6139122
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
@@ -1139,7 +1139,7 @@ def get_module_duration(module):
     duration = module.__duration__
     if duration == 0:  # e.g. ModuleList
         for m in module.children():
-            duration += m.__duration__
+            duration += get_module_duration(m)
     return duration