code: def train_model(
      model, criterion, optimizer, scheduler, num_epochs, data_dir
  ):
      """Train the model."""
  
      # load training/validation data
      dataloaders, dataset_sizes, class_names = load_data(data_dir)
  
      device = torch.device(
          "cuda:0" if torch.cuda.is_available() else "cpu"
      )
  
      since = time.time()
  
      best_model_wts = copy.deepcopy(model.state_dict())
      best_acc = 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", "val"]:
              if phase == "train":
                  scheduler.step()
                  model.train()  # Set model to training mode
              else:
                  model.eval()  # Set model to evaluate mode
  
              running_loss = 0.0
              running_corrects = 0
  
              # Iterate over data.
              for inputs, labels in dataloaders[phase]:
                  inputs = inputs.to(device)
                  labels = labels.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, labels)
  
                      # backward + optimize only if in training phase
                      if phase == "train":
                          loss.backward()
                          optimizer.step()
  
                  # statistics
                  running_loss += loss.item() * inputs.size(0)
                  running_corrects += torch.sum(preds == labels.data)
  
              epoch_loss = running_loss / dataset_sizes[phase]
              epoch_acc = (
                  running_corrects.double() / dataset_sizes[phase]
              )
  
              print(
                  "{} Loss: {:.4f} Acc: {:.4f}".format(
                      phase, epoch_loss, epoch_acc
                  )
              )
  
              # deep copy the model
              if phase == "val" and epoch_acc > best_acc:
                  best_acc = epoch_acc
                  best_model_wts = copy.deepcopy(model.state_dict())
  
              # log the best val accuracy to AML run
              run.log("best_val_acc", np.float(best_acc))
  
          print()
  
      time_elapsed = time.time() - since
      print(
          "Training complete in {:.0f}m {:.0f}s".format(
              time_elapsed // 60, time_elapsed % 60
          )
      )
      print("Best val Acc: {:4f}".format(best_acc))
  
      # load best model weights
      model.load_state_dict(best_model_wts)
      return model
  
file_path: MachineLearningNotebooks/how-to-use-azureml/ml-frameworks/pytorch/train-hyperparameter-tune-deploy-with-pytorch/pytorch_train.py
imports: from __future__ import print_function, division
  import torch
  import torch.nn
  import torch.optim
  from torch.optim import lr_scheduler
  from torchvision import datasets, models, transforms
  import numpy
  import time
  import os
  import copy
  import argparse
  from azureml.core.run import Run
label_desc: Use image node_id 2 for calling a global function with example usage: train_model(model, criterion, optimizer, scheduler, num_epochs, data_dir) and returns: model