class_docstr: ## Retro trainer
class_docstr_tok: ["#","#","Retro","trainer"]
docstr: ### Train the model for an epoch
docstr_tok: ["#","#","#","Train","the","model","for","an","epoch"]
code: def __call__(self):
      """
      ### Train the model for an epoch
      """
  
      # Iterate through training data
      for i, (src, tgt, neighbors) in monit.enum(
          "Train", self.dataloader
      ):
          # Move data to the device
          src, tgt, neighbors = (
              src.to(self.device),
              tgt.to(self.device),
              neighbors.to(self.device),
          )
  
          # Forward pass
          res = self.model(src, neighbors)
          # Calculate loss
          loss = self.loss_func(
              res.view(-1, res.shape[-1]), tgt.view(-1)
          )
  
          # Clear the gradients
          self.optimizer.zero_grad()
          # Backward pass
          loss.backward()
          # Optimize the model
          self.optimizer.step()
  
          # Save training statistics and increment the global step counter
          tracker.save({"loss.train": loss})
          tracker.add_global_step(len(src))
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/retro/train.py
imports: import torch
  from torch import nn
  from torch.utils.data import DataLoader, RandomSampler
  from labml import monit, lab, tracker, experiment, logger
  from labml.logger import Text
  from labml_helpers.datasets.text import TextFileDataset
  from labml_nn.optimizers.noam import Noam
  from labml_nn.transformers.retro import model
  from labml_nn.transformers.retro.dataset import Dataset, RetroIndex
  from labml_nn.transformers.retro.model import RetroModel, NearestNeighborEncoder
label_desc: Use image node_id 2 for calling the Trainer obj's underlying member method code with example usage: obj.__call__() without return types