code: def trainModel(model, loader, optimizer, device, criterion=F.nll_loss,
                 batches_in_epoch=sys.maxsize, batch_callback=None,
                 progress_bar=None):
    
    model.train()
    if progress_bar is not None:
      loader = tqdm(loader, **progress_bar)
      # update progress bar total based on batches_in_epoch
      if batches_in_epoch < len(loader):
        loader.total = batches_in_epoch
  
    for batch_idx, (data, target) in enumerate(loader):
      data, target = data.to(device), target.to(device)
      optimizer.zero_grad()
      output = model(data)
      loss = criterion(output, target)
      loss.backward()
      optimizer.step()
  
      if batch_callback is not None:
        batch_callback(model=model, batch_idx=batch_idx)
      if batch_idx >= batches_in_epoch:
        break
  
    if progress_bar is not None:
      loader.n = loader.total
      loader.close()
docs: Train the given model by iterating through mini batches. An epoch
    ends after one pass through the training set, or if the number of mini
    batches exceeds the parameter "batches_in_epoch".
  
    :param model: pytorch model to be trained
    :type model: torch.nn.Module
    :param loader: train dataset loader
    :type loader: :class:`torch.utils.data.DataLoader`
    :param optimizer: Optimizer object used to train the model.
           This function will train the model on every batch using this optimizer
           and the :func:`torch.nn.functional.nll_loss` function
    :param batches_in_epoch: Max number of mini batches to train.
    :param device: device to use ('cpu' or 'cuda')
    :type device: :class:`torch.device
    :param criterion: loss function to use
    :type criterion: function
    :param batch_callback: Callback function to be called on every batch with the
                           following parameters: model, batch_idx
    :type batch_callback: function
    :param progress_bar: Optional :class:`tqdm` progress bar args.
                         None for no progress bar
    :type progress_bar: dict or None