def on_start(self, task: "tasks.ClassyTask") -> None:
"""
Logs Gpu nvidia-smi stats to logger streams.
"""
if is_primary() and (task.device.type == "cuda"):
# print the nvidia-smi stats
log_gpu_stats()
def on_step(self, task: "tasks.ClassyTask") -> None:
"""
Print the nvidia-smi stats again to get more accurate nvidia-smi
useful for monitoring memory usage.
"""
if (is_primary() and (task.device.type == "cuda")
and task.local_iteration_num == 50):
log_gpu_stats()
