fix: remove all buffers registered beside the supported ops

Lyken17 · Mar 9, 2023 · ae677a0 · ae677a0
1 parent 43c064a
commit ae677a0
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/thop/profile.py b/thop/profile.py
@@ -203,6 +203,10 @@ def add_hooks(m: nn.Module):
             )
         types_collection.add(m_type)
 
+    def remove_buffers(m: nn.Module):
+        m._buffers.pop("total_ops")
+        m._buffers.pop("total_params")
+
     prev_training_status = model.training
 
     model.eval()
@@ -239,8 +243,7 @@ def dfs_count(module: nn.Module, prefix="\t") -> (int, int):
     for m, (op_handler, params_handler) in handler_collection.items():
         op_handler.remove()
         params_handler.remove()
-        m._buffers.pop("total_ops")
-        m._buffers.pop("total_params")
+    model.apply(remove_buffers)
 
     if ret_layer_info:
         return total_ops, total_params, ret_dict