remove profiler

This commit is contained in:
wfjsw 2024-02-12 00:25:40 -06:00
parent a58ee39e38
commit cf3cc4c762

View File

@ -208,7 +208,6 @@ class ModelMover:
handle.remove()
def _pre_forward_hook(self, module, _):
with torch.profiler.record_function("lowvram prehook"):
with stream_wrapper(stream=self.model_mover_stream):
idx = self.submodules_indexer[module]
for i in range(idx, idx + self.lookahead_distance):
@ -231,7 +230,6 @@ class ModelMover:
submodule.to(devices.device, non_blocking=True)
def _post_forward_hook(self, module, _1, _2):
with torch.profiler.record_function("lowvram posthook"):
with stream_wrapper(stream=self.model_mover_stream):
del self.module_movement_events[module]
module.to(cpu, non_blocking=True)