code: def main():
      epochs = 10
  
      is_save_models = True
      train_batch_size = 64
      valid_batch_size = 1000
  
      use_cuda = True
      seed = 5
      train_log_interval = 10
  
      learning_rate = 0.01
  
      # get device
      is_cuda = use_cuda and torch.cuda.is_available()
      if not is_cuda:
          device = torch.device("cpu")
      else:
          device = torch.device(f"cuda:0")
  
      # data transform
      data_transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize((0.1307,), (0.3081,)),
          ]
      )
  
      # train loader
      train_loader = torch.utils.data.DataLoader(
          datasets.MNIST(
              str(lab.get_data_path()),
              train=True,
              download=True,
              transform=data_transform,
          ),
          batch_size=train_batch_size,
          shuffle=True,
      )
  
      # valid loader
      valid_loader = torch.utils.data.DataLoader(
          datasets.MNIST(
              str(lab.get_data_path()),
              train=False,
              download=True,
              transform=data_transform,
          ),
          batch_size=valid_batch_size,
          shuffle=False,
      )
  
      # model
      model = Net().to(device)
  
      # optimizer
      optimizer = optim.Adam(model.parameters(), lr=learning_rate)
  
      # set seeds
      torch.manual_seed(seed)
  
      # training loop
      for epoch in range(1, epochs + 1):
          train(
              epoch,
              model,
              optimizer,
              train_loader,
              device,
              train_log_interval,
          )
          validate(epoch, model, valid_loader, device)
  
      if is_save_models:
          torch.save(model.state_dict(), "mnist_cnn.pt")
  
file_path: labml/samples/pytorch/mnist/a_no_labml.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torch.utils.data
  from torchvision import datasets, transforms
  from labml import lab
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types