code: def run(hparams):
      torch.manual_seed(hparams.seed)
  
      use_cuda = torch.cuda.is_available()
      use_mps = torch.backends.mps.is_available()
      if use_cuda:
          device = torch.device("cuda")
      elif use_mps:
          device = torch.device("mps")
      else:
          device = torch.device("cpu")
  
      transform = T.Compose(
          [T.ToTensor(), T.Normalize((0.1307,), (0.3081,))]
      )
      train_dataset = MNIST(
          DATASETS_PATH, train=True, download=True, transform=transform
      )
      test_dataset = MNIST(
          DATASETS_PATH, train=False, transform=transform
      )
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=hparams.batch_size,
      )
      test_loader = torch.utils.data.DataLoader(
          test_dataset, batch_size=hparams.batch_size
      )
  
      model = Net().to(device)
      optimizer = optim.Adadelta(model.parameters(), lr=hparams.lr)
  
      scheduler = StepLR(optimizer, step_size=1, gamma=hparams.gamma)
  
      # EPOCH LOOP
      for epoch in range(1, hparams.epochs + 1):
          # TRAINING LOOP
          model.train()
          for batch_idx, (data, target) in enumerate(train_loader):
              data, target = data.to(device), target.to(device)
              optimizer.zero_grad()
              output = model(data)
              loss = F.nll_loss(output, target)
              loss.backward()
              optimizer.step()
              if (batch_idx == 0) or (
                  (batch_idx + 1) % hparams.log_interval == 0
              ):
                  print(
                      "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format(
                          epoch,
                          batch_idx * len(data),
                          len(train_loader.dataset),
                          100.0 * batch_idx / len(train_loader),
                          loss.item(),
                      )
                  )
                  if hparams.dry_run:
                      break
          scheduler.step()
  
          # TESTING LOOP
          model.eval()
          test_loss = 0
          correct = 0
          with torch.no_grad():
              for data, target in test_loader:
                  data, target = data.to(device), target.to(device)
                  output = model(data)
                  test_loss += F.nll_loss(
                      output, target, reduction="sum"
                  ).item()  # sum up batch loss
                  pred = output.argmax(
                      dim=1, keepdim=True
                  )  # get the index of the max log-probability
                  correct += pred.eq(target.view_as(pred)).sum().item()
                  if hparams.dry_run:
                      break
  
          test_loss /= len(test_loader.dataset)
  
          print(
              "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(
                  test_loss,
                  correct,
                  len(test_loader.dataset),
                  100.0 * correct / len(test_loader.dataset),
              )
          )
  
          if hparams.dry_run:
              break
  
      if hparams.save_model:
          torch.save(model.state_dict(), "mnist_cnn.pt")
  
file_path: pytorch-lightning/examples/fabric/image_classifier/train_torch.py
imports: import argparse
  from os import path
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torchvision.transforms
  from torch.optim.lr_scheduler import StepLR
  from torchvision.datasets import MNIST
label_desc: Use image node_id 1 for calling a global function with example usage: run(hparams) without return types