| 278 | |
| 279 | |
| 280 | def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor: |
| 281 | if isinstance(parameters, torch.Tensor): |
| 282 | parameters = [parameters] |
| 283 | parameters = [p for p in parameters if p.grad is not None] |
| 284 | norm_type = float(norm_type) |
| 285 | if len(parameters) == 0: |
| 286 | return torch.tensor(0.) |
| 287 | device = parameters[0].grad.device |
| 288 | if norm_type == inf: |
| 289 | total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters) |
| 290 | else: |
| 291 | total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type) |
| 292 | return total_norm |
| 293 | |
| 294 | |
| 295 | def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler): |