MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / clear_state

Method clear_state

tensorrt_llm/_torch/debug/debug_hook.py:59–80  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

57 self.is_pre_forward = is_pre_forward
58
59 def clear_state(self):
60 self.pre_forward_actions.clear()
61 self.after_forward_actions.clear()
62 self.layer_names.clear()
63 self.layer_inner_counter.clear()
64
65 if self.module_forward_hook_handle is not None:
66 self.module_forward_hook_handle.remove()
67 if self.module_forward_pre_hook_handle is not None:
68 self.module_forward_pre_hook_handle.remove()
69
70 self.module_forward_hook_handle = None
71 self.module_forward_pre_hook_handle = None
72
73 for _, handler in self.forward_hook_handles.items():
74 handler.remove()
75
76 for _, handler in self.forward_pre_hook_handles.items():
77 handler.remove()
78
79 self.forward_hook_handles.clear()
80 self.forward_pre_hook_handles.clear()
81
82 def register_pre_forward_action(self, filter, action):
83 self.pre_forward_actions.append((filter, action))

Callers 1

disable_debugFunction · 0.80

Calls 2

clearMethod · 0.45
removeMethod · 0.45

Tested by

no test coverage detected