Thanks, here is my code that runs run_epoch. I did not include this part because it takes inputs from my custom made dataloaders (a feature which I excluded from the minimal model), and other than that all it does is loop (but I should have, since I just noticed the optimizer is here).  Besides, are you going to make minimal inputs yourself or do you need the images that I used for the model? Really appreciate your help.
def run_model(model,src_batch, tar_batch, grid_size=8,):
  optimizer = optim.Adam(model.parameters(), lr=1e-3)
  epoch_loss = []
  for i in range(10):
    loss_list = run_epoch(model, optimizer, src_batch, tar_batch , grid_size)
    avg_epoch_loss = sum(loss_list) / len(loss_list)
    print(f'Loss in Epoch {i}: {avg_epoch_loss}')