modelzoo.common.pytorch.run_cstorch_flow.optimizer_step_with_summaries#

modelzoo.common.pytorch.run_cstorch_flow.optimizer_step_with_summaries(loss: torch.Tensor, optimizer: cerebras_pytorch.optim.optimizer.Optimizer, grad_scaler: cerebras_pytorch.amp.grad_scaler.GradScaler, max_gradient_norm: Optional[float] = None, max_gradient_value: Optional[float] = None, log_summaries: bool = False, model: Optional[torch.nn.Module] = None)[source]#

Customized equivalent to cstorch.amp.optimizer_step additionally featuring grad norm summaries

previous

modelzoo.common.pytorch.run_cstorch_flow.log_input_summary

next

modelzoo.common.pytorch.run_cstorch_flow.run_cstorch_eval