| 342 | return control |
| 343 | |
| 344 | class ProfilerTrainer(Trainer): |
| 345 | def __init__(self, *args, **kwargs): |
| 346 | super().__init__(*args, **kwargs) |
| 347 | self.profiler = torch.profiler.profile( |
| 348 | schedule=torch.profiler.schedule(wait=2, warmup=2, active=4), |
| 349 | on_trace_ready=torch.profiler.tensorboard_trace_handler("./profiler_output") |
| 350 | ) |
| 351 | self.profiler.__enter__() |
| 352 | |
| 353 | def training_step(self, model, inputs): |
| 354 | output = super().training_step(model, inputs) |
| 355 | self.profiler.step() |
| 356 | return output |
| 357 | |
| 358 | def __del__(self): |
| 359 | self.profiler.__exit__(None, None, None) |
nothing calls this directly
no outgoing calls
no test coverage detected