text: def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
      since = time.time()
  
      best_model_wts = copy.deepcopy(model.state_dict())
      best_loss = 0.0
  
      for epoch in range(num_epochs):
          print('Epoch {}/{}'.format(epoch, num_epochs - 1))
          print('-' * 10)
  
          # Each epoch has a training and validation phase
          for phase in ['train']:
              if phase == 'train':
                  model.train()  # Set model to training mode
              else:
                  model.eval()   # Set model to evaluate mode
  
              running_loss = 0.0
              
  
              # Iterate over data.
              for data in dataloader:
                  inputs, params = data['image'], data['params']
                  inputs = inputs.to(device)
                  params = params.to(device)
  
                  # zero the parameter gradients
                  optimizer.zero_grad()
  
                  # forward
                  # track history if only in train
                  with torch.set_grad_enabled(phase == 'train'):
                      outputs = model(inputs)
                      _, preds = torch.max(outputs, 1)
                      loss = criterion(outputs, params)
  
                      # backward + optimize only if in training phase
                      if phase == 'train':
                          loss.backward()
                          optimizer.step()
  
                  # statistics
                  running_loss += loss.item() * inputs.size(0)
              if phase == 'train':
                  scheduler.step()
  
              epoch_loss = running_loss / len(mnu_dataset)
              
              print('{} Loss: {:.4f} '.format(
                  phase, epoch_loss))
  
              # # deep copy the model
              # if phase == 'val' and epoch_loss < best_loss:
              #     best_loss = epoch_loss
              #     best_model_wts = copy.deepcopy(model.state_dict())
  
          print()
  
      time_elapsed = time.time() - since
      print('Training complete in {:.0f}m {:.0f}s'.format(
          time_elapsed // 60, time_elapsed % 60))
      print('Best val Loss: {:4f}'.format(best_loss))
  
      # load best model weights
      # model.load_state_dict(best_model_wts)
      return model