text: import torch
  import numpy as np
  from torchvision import datasets
  import torchvision.transforms as transforms
  from torch.utils.data.sampler import SubsetRandomSampler
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  import matplotlib.pyplot as plt
  from network import ConvCIFAR
  from ada_hessian import AdaHessian
  import time
  
  # list of all optimizers
  optimizer_type = "AdaHessian"
  # optimizer_type = "SGD"
  # optimizer_type = "Adam"
  # optimizer_type = "Adagrad"
  # optimizer_type = "AdamW"
  # optimizer_type = "RMSProp"
  
  
  # Check if CUDA is available
  train_on_gpu = torch.cuda.is_available()
  
  if train_on_gpu:
      print("CUDA is available! Training on GPU...")
  else:
      print("CUDA is not available. Training on CPU...")
  
  # Number of subprocesses to use for data loading
  num_workers = 0
  
  # How many samples per batch to load
  batch_size = 20
  
  # Percentage of training set to use as validation
  n_valid = 0.2
  
  # Convert data to a normalized torch.FloatTensor
  # Data augmentation
  transform = transforms.Compose(
      [
          transforms.RandomHorizontalFlip(),  # randomly flip and rotate
          transforms.RandomRotation(10),
          transforms.ToTensor(),
          transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
      ]
  )
  
  # Select training_set and testing_set
  train_data = datasets.CIFAR10("../data", train=True, download=True, transform=transform)
  
  test_data = datasets.CIFAR10("../data", train=False, download=True, transform=transform)
  
  # Get indices for training_set and validation_set
  n_train = len(train_data)
  indices = list(range(n_train))
  np.random.shuffle(indices)
  split = int(np.floor(n_valid * n_train))
  train_idx, valid_idx = indices[split:], indices[:split]
  
  # Define samplers for obtaining training and validation
  train_sampler = SubsetRandomSampler(train_idx)
  valid_sampler = SubsetRandomSampler(valid_idx)
  
  # Prepare data loaders (combine dataset and sampler)
  train_loader = torch.utils.data.DataLoader(
      train_data, batch_size=batch_size, sampler=train_sampler, num_workers=num_workers
  )
  
  valid_loader = torch.utils.data.DataLoader(
      train_data, batch_size=batch_size, sampler=valid_sampler, num_workers=num_workers
  )
  
  test_loader = torch.utils.data.DataLoader(
      test_data, batch_size=batch_size, num_workers=num_workers
  )
  
  # Specify the image classes
  classes = [
      "airplane",
      "automobile",
      "bird",
      "cat",
      "deer",
      "dog",
      "frog",
      "horse",
      "ship",
      "truck",
  ]
  
  model = ConvCIFAR()
  
  if train_on_gpu:
      model.cuda()
  
  # Specify the Loss function
  criterion = nn.CrossEntropyLoss()
  
  # Specify the optimizer
  if optimizer_type == "AdaHessian":
      optimizer = AdaHessian(model.parameters())
  elif optimizer_type == "SGD":
      optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
  elif optimizer_type == "Adam":
      optimizer = optim.Adam(model.parameters(), lr=0.001)
  elif optimizer_type == "Adagrad":
      optimizer = optim.Adagrad(model.parameters(), lr=0.01)
  elif optimizer_type == "AdamW":
      optimizer = optim.AdamW(model.parameters(), lr=0.001)
  elif optimizer_type == "RMSProp":
      optimizer = optim.RMSprop(model.parameters(), lr=0.01)
  else:
      raise TypeError(f"invalid optimizer type: {optimizer_type}")
  
  # number of epochs to train the model
  n_epochs = 5
  
  valid_loss_min = np.Inf  # track change in validation loss
  
  for epoch in range(1, n_epochs + 1):
  
      # keep track of training and validation loss
      train_loss = 0.0
      valid_loss = 0.0
  
      # train the model #
      model.train()
      start = time.time()
      for data, target in train_loader:
          # move tensors to GPU if CUDA is available
  
          if train_on_gpu:
              data, target = data.cuda(), target.cuda()
          # clear the gradients of all optimized variables
          optimizer.zero_grad()
          # forward pass: compute predicted outputs by passing inputs to the model
          output = model(data)
          # calculate the batch loss
          loss = criterion(output, target)
          # backward pass: compute gradient of the loss with respect to model parameters
          if optimizer_type == "AdaHessian":
              loss.backward(create_graph=True)
          else:
              loss.backward()
          # perform a single optimization step (parameter update)
          optimizer.step()
          # update training loss
          train_loss += loss.item() * data.size(0)
      print(f"Loss: {train_loss}\tIteration Cost: {round(time.time() - start, 3)}")
  
      # validate the model #
      model.eval()
      for data, target in valid_loader:
          # move tensors to GPU if CUDA is available
          if train_on_gpu:
              data, target = data.cuda(), target.cuda()
          # forward pass: compute predicted outputs by passing inputs to the model
          output = model(data)
          # calculate the batch loss
          loss = criterion(output, target)
          # update average validation loss
          valid_loss += loss.item() * data.size(0)
  
      # calculate average losses
      train_loss = train_loss / len(train_loader.dataset)
      valid_loss = valid_loss / len(valid_loader.dataset)
  
      # print training/validation statistics
      print(
          "Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}".format(
              epoch, train_loss, valid_loss
          )
      )
  
      # save model if validation loss has decreased
      if valid_loss <= valid_loss_min:
          print(
              "Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...".format(
                  valid_loss_min, valid_loss
              )
          )
          torch.save(model.state_dict(), "../model/cnn_cifar.pt")
          valid_loss_min = valid_loss
  
      # evaluate model
      model.load_state_dict(torch.load("../model/cnn_cifar.pt"))
  
      # track test loss
      test_loss = 0.0
      class_correct = list(0.0 for i in range(10))
      class_total = list(0.0 for i in range(10))
  
      model.eval()
      # iterate over test data
      for data, target in test_loader:
          # move tensors to GPU if CUDA is available
          if train_on_gpu:
              data, target = data.cuda(), target.cuda()
          # forward pass: compute predicted outputs by passing inputs to the model
          output = model(data)
          # calculate the batch loss
          loss = criterion(output, target)
          # update test loss
          test_loss += loss.item() * data.size(0)
          # convert output probabilities to predicted class
          _, pred = torch.max(output, 1)
          # compare predictions to true label
          correct_tensor = pred.eq(target.data.view_as(pred))
          correct = (
              np.squeeze(correct_tensor.numpy())
              if not train_on_gpu
              else np.squeeze(correct_tensor.cpu().numpy())
          )
          # calculate test accuracy for each object class
          for i in range(batch_size):
              label = target.data[i]
              class_correct[label] += correct[i].item()
              class_total[label] += 1
  
      # average test loss
      test_loss = test_loss / len(test_loader.dataset)
      print("Test Loss: {:.6f}\n".format(test_loss))
  
      for i in range(10):
          if class_total[i] > 0:
              print(
                  "Test Accuracy of %5s: %2d%% (%2d/%2d)"
                  % (
                      classes[i],
                      100 * class_correct[i] / class_total[i],
                      np.sum(class_correct[i]),
                      np.sum(class_total[i]),
                  )
              )
          else:
              print("Test Accuracy of %5s: N/A (no training examples)" % (classes[i]))
  
      print(
          "\nTest Accuracy (Overall): %2d%% (%2d/%2d)"
          % (
              100.0 * np.sum(class_correct) / np.sum(class_total),
              np.sum(class_correct),
              np.sum(class_total),
          )
      )
  