code: def train_one_epoch(
      model, optimizer, scheduler, train_loader, logger, args
  ):
      device = torch.device(args.device)
      for data_blob in logger.log_every(train_loader):
          optimizer.zero_grad()
  
          image1, image2, flow_gt, valid_flow_mask = (
              x.to(device) for x in data_blob
          )
          flow_predictions = model(
              image1, image2, num_flow_updates=args.num_flow_updates
          )
  
          loss = utils.sequence_loss(
              flow_predictions, flow_gt, valid_flow_mask, args.gamma
          )
          metrics, _ = utils.compute_metrics(
              flow_predictions[-1], flow_gt, valid_flow_mask
          )
  
          metrics.pop("f1")
          logger.update(loss=loss, **metrics)
  
          loss.backward()
  
          torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1)
  
          optimizer.step()
          scheduler.step()
  
file_path: vision/references/optical_flow/train.py
imports: import argparse
  import warnings
  from math import ceil
  from pathlib import Path
  import torch
  import torchvision.models.optical_flow
  import utils
  from presets import OpticalFlowPresetEval, OpticalFlowPresetTrain
  from torchvision.datasets import FlyingChairs, FlyingThings3D, HD1K, KittiFlow, Sintel
label_desc: Use image node_id 4 for calling a global function with example usage: train_one_epoch(model, optimizer, scheduler, train_loader, logger, args) without return types