Conditionally deactivates the profiler based on the current step. This method checks if the current training step matches the step designated for finishing the initial profile, or if it meets the criteria for deactivating a periodic profile.
(self, step, state)
| 88 | jax.profiler.start_trace(self.output_path) |
| 89 | |
| 90 | def maybe_deactivate_profiler(self, step, state): |
| 91 | """Conditionally deactivates the profiler based on the current step. |
| 92 | This method checks if the current training step matches the step designated |
| 93 | for finishing the initial profile, or if it meets the criteria for |
| 94 | deactivating a periodic profile. |
| 95 | """ |
| 96 | if self.mode != "" and (step == self.finished_initial_profile_step or self.should_deactivate_periodic_profile(step)): |
| 97 | self.deactivate(blocking_object=state) |
| 98 | |
| 99 | def deactivate(self, blocking_object=None): |
| 100 | """End the profiler. |
no test coverage detected