cerebras.modelzoo.common.run_cstorch_flow.optimizer_step_with_summaries#

cerebras.modelzoo.common.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