Thanks for the response. I made a minimal version of  my function that run a single epoch. Please let me know if you need anything else. Thanks!
#aug_batch, src_batch, tar_batch input form: nx128x128
#run the model for a single epoch
#data_loader can either be a single data_loader or a iterable list of data loaders (if so, source loader should also be a list of data loaders)
def run_epoch(model, optimizer,  aug_batch, src_batch, tar_batch, grid_size=8):

      aug_batch = torch.tensor(aug_batch, dtype=torch.float32)
      tar_batch = torch.tensor(tar_batch, dtype=torch.float32)
      src_batch = torch.tensor(src_batch, dtype=torch.float32)

      #run forward propagation
      tar_est = warp(model, src_batch, aug_batch, grid_size)        
      tar_est = tar_est.squeeze(dim=1)
      total_loss = Total_Loss(tar_batch, tar_est, grid_size, 1e-5)

      total_loss.backward()
      optimizer.step()
      optimizer.zero_grad()
      loss_list.append(total_loss)

  return loss_list


def warp(model, src_batch, aug_batch, grid_size):
  input_image = torch.stack([src_batch, aug_batch])
  input_image  = input_image.permute([1,0,2,3])
  #run the network
  target_est = model.forward(input_image, src_batch)

  return target_est```