code: def train_one_epoch(
      model, criterion, optimizer, data_loader, device, epoch
  ):
      model.train()
  
      sums = defaultdict(lambda: 0.0)
      start1 = time()
  
      metric = MetricLogger("train_iteration")
      metric["epoch"] = epoch
  
      for waveform, specgram, target in data_loader:
          start2 = time()
  
          waveform = waveform.to(device)
          specgram = specgram.to(device)
          target = target.to(device)
  
          output = model(waveform, specgram)
          output, target = output.squeeze(1), target.squeeze(1)
  
          loss = criterion(output, target)
          loss_item = loss.item()
          sums["loss"] += loss_item
          metric["loss"] = loss_item
  
          optimizer.zero_grad()
          loss.backward()
  
          if args.clip_grad > 0:
              gradient = torch.nn.utils.clip_grad_norm_(
                  model.parameters(), args.clip_grad
              )
              sums["gradient"] += gradient.item()
              metric["gradient"] = gradient.item()
  
          optimizer.step()
  
          metric["iteration"] = sums["iteration"]
          metric["time"] = time() - start2
          metric()
          sums["iteration"] += 1
  
      avg_loss = sums["loss"] / len(data_loader)
  
      metric = MetricLogger("train_epoch")
      metric["epoch"] = epoch
      metric["loss"] = sums["loss"] / len(data_loader)
      metric["gradient"] = avg_loss
      metric["time"] = time() - start1
      metric()
  
file_path: audio/examples/pipeline_wavernn/main.py
imports: import argparse
  import logging
  import os
  from collections import defaultdict
  from datetime import datetime
  from time import time
  from typing import List
  import torch
  import torchaudio
  from datasets import collate_factory, split_process_dataset
  from losses import LongCrossEntropyLoss, MoLLoss
  from processing import NormalizeDB
  from torch.optim import Adam
  from torch.utils.data import DataLoader
  from torchaudio.models.wavernn import WaveRNN
  from utils import count_parameters, MetricLogger, save_checkpoint
label_desc: Use image node_id 2 for calling a global function with example usage: train_one_epoch(model, criterion, optimizer, data_loader, device, epoch) without return types