code: def _train_one_epoch(
      model: Union[DDP, DMP],
      train_loader: data_utils.DataLoader,
      device: torch.device,
      optimizer: optim.Adam,
      lr_scheduler: optim.lr_scheduler.StepLR,
      epoch: int,
  ) -> None:
      """
      Train model for 1 epoch. Helper function for train_val_test.
  
      Args:
          model (Union[DDP, DMP]): DMP or DDP model contains the Bert4Rec.
          train_loader (data_utils.DataLoader): DataLoader used for training.
          device (torch.device): the device to train/val/test
          optimizer (optim.Adam): Adam optimizer to train the model
          lr_scheduler (optim.lr_scheduler.StepLR): scheduler to control the learning rate
          epoch (int): the current epoch number
  
      Returns:
          None.
      """
      model.train()
      if torch.cuda.is_available():
          torch.cuda.set_device(dist.get_rank())
      loss_logs = []
      train_iterator = iter(train_loader)
      ce = nn.CrossEntropyLoss(ignore_index=0)
      outputs = [None for _ in range(dist.get_world_size())]
      for _ in tqdm(iter(int, 1), desc=f"Epoch {epoch+1}"):
          try:
              batch = next(train_iterator)
              batch = [x.to(device) for x in batch]
  
              optimizer.zero_grad()
              seqs, labels = batch
  
              kjt = _to_kjt(seqs, device)
              logits = model(kjt)  # B x T x V
  
              logits = logits.view(-1, logits.size(-1))  # (B*T) x V
              labels = labels.view(-1)  # B*T
              loss = ce(logits, labels)
  
              loss.backward()
  
              optimizer.step()
  
              loss_logs.append(loss.item())
  
          except StopIteration:
              break
      dist.all_gather_object(outputs, sum(loss_logs) / len(loss_logs))
      if dist.get_rank() == 0:
          # pyre-fixme[6]: For 1st param expected `Iterable[Variable[_SumT (bound to
          #  _SupportsSum)]]` but got `List[None]`.
          print(
              f"Epoch {epoch + 1}, average loss { (sum(outputs) or 0) /len(outputs)}"
          )
      lr_scheduler.step()
  
imports: import argparse
  import os
  import sys
  from typing import Any, cast, Dict, List, Union
  import numpy
  import torch
  import torch.nn
  import torch.optim
  import torch.utils.data
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from torch import distributed
  from torch.nn.parallel import DistributedDataParallel
  from torchrec.distributed.embedding import EmbeddingCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.types import ModuleSharder
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from tqdm import tqdm
label_desc: Use image node_id 5 for calling a global function with example usage: _train_one_epoch(model, train_loader, device, optimizer, lr_scheduler, epoch) without return types