code: def train(
      net: torch.nn.Module,
      train_loader: DataLoader,
      parameters: Dict[str, float],
      dtype: torch.dtype,
      device: torch.device,
  ) -> nn.Module:
      """
      Train CNN on provided data set.
  
      Args:
          net: initialized neural network
          train_loader: DataLoader containing training set
          parameters: dictionary containing parameters to be passed to the optimizer.
              - lr: default (0.001)
              - momentum: default (0.0)
              - weight_decay: default (0.0)
              - num_epochs: default (1)
          dtype: torch dtype
          device: torch device
      Returns:
          nn.Module: trained CNN.
      """
      # Initialize network
      net.to(dtype=dtype, device=device)
      net.train()
      # Define loss and optimizer
      criterion = nn.NLLLoss(reduction="sum")
      optimizer = optim.SGD(
          net.parameters(),
          lr=parameters.get("lr", 0.001),
          momentum=parameters.get("momentum", 0.0),
          weight_decay=parameters.get("weight_decay", 0.0),
      )
      scheduler = optim.lr_scheduler.StepLR(
          optimizer,
          step_size=int(parameters.get("step_size", 30)),
          gamma=parameters.get(
              "gamma", 1.0
          ),  # default is no learning rate decay
      )
      num_epochs = parameters.get("num_epochs", 1)
  
      # Train Network
      # pyre-fixme[6]: Expected `int` for 1st param but got `float`.
      for _ in range(num_epochs):
          for inputs, labels in train_loader:
              # move data to proper dtype and device
              inputs = inputs.to(dtype=dtype, device=device)
              labels = labels.to(device=device)
  
              # zero the parameter gradients
              optimizer.zero_grad()
  
              # forward + backward + optimize
              outputs = net(inputs)
              loss = criterion(outputs, labels)
              loss.backward()
              optimizer.step()
              scheduler.step()
      return net
  
file_path: Ax/ax/utils/tutorials/cnn_utils.py
imports: from itertools import accumulate
  from typing import Dict, List, Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torchvision
  import torchvision.transforms
  from torch._tensor import Tensor
  from torch.utils.data import DataLoader, Dataset, Subset
label_desc: Use image node_id 4 for calling a global function with example usage: train(net, train_loader, parameters, dtype, device) and returns: net