class_docstr: ## Configurations
class_docstr_tok: ["#","#","Configurations"]
docstr: ### Train for an epoch
docstr_tok: ["#","#","#","Train","for","an","epoch"]
code: def train(self):
      """
      ### Train for an epoch
      """
  
      # Iterate through the dataset.
      # Use [`mix`](https://docs.labml.ai/api/monit.html#labml.monit.mix)
      # to sample $50$ times per epoch.
      for _, (image, mask) in monit.mix(
          ("Train", self.data_loader), (self.sample, list(range(50)))
      ):
          # Increment global step
          tracker.add_global_step()
          # Move data to device
          image, mask = image.to(self.device), mask.to(self.device)
  
          # Make the gradients zero
          self.optimizer.zero_grad()
          # Get predicted mask logits
          logits = self.model(image)
          # Crop the target mask to the size of the logits. Size of the logits will be smaller if we
          # don't use padding in convolutional layers in the U-Net.
          mask = torchvision.transforms.functional.center_crop(
              mask, [logits.shape[2], logits.shape[3]]
          )
          # Calculate loss
          loss = self.loss_func(self.sigmoid(logits), mask)
          # Compute gradients
          loss.backward()
          # Take an optimization step
          self.optimizer.step()
          # Track the loss
          tracker.save("loss", loss)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/unet/experiment.py
imports: import numpy
  import torch
  import torch.utils.data
  import torchvision.transforms.functional
  from torch import nn
  from labml import lab, tracker, experiment, monit
  from labml.configs import BaseConfigs
  from labml_helpers.device import DeviceConfigs
  from labml_nn.unet.carvana import CarvanaDataset
  from labml_nn.unet import UNet
label_desc: Use image node_id 3 for calling the Configs obj's underlying member method code with example usage: obj.train() without return types