code: def main():
      conf = Config()
      log_dir = "runs/mnist_custom_optim"
      print("Tensorboard: tensorboard --logdir={}".format(log_dir))
  
      with SummaryWriter(log_dir) as writer:
          use_cuda = not conf.no_cuda and torch.cuda.is_available()
          torch.manual_seed(conf.seed)
          device = torch.device("cuda" if use_cuda else "cpu")
          train_loader, test_loader = prepare_loaders(conf, use_cuda)
  
          model = Net().to(device)
  
          # create grid of images and write to tensorboard
          images, labels = next(iter(train_loader))
          img_grid = utils.make_grid(images)
          writer.add_image("mnist_images", img_grid)
  
          # custom optimizer from torch_optimizer package
          optimizer = optim.DiffGrad(model.parameters(), lr=conf.lr)
  
          scheduler = StepLR(optimizer, step_size=1, gamma=conf.gamma)
          for epoch in range(1, conf.epochs + 1):
              train(
                  conf,
                  model,
                  device,
                  train_loader,
                  optimizer,
                  epoch,
                  writer,
              )
              test(conf, model, device, test_loader, epoch, writer)
              scheduler.step()
              for name, param in model.named_parameters():
                  writer.add_histogram(name, param, epoch)
                  writer.add_histogram(
                      "{}.grad".format(name), param.grad, epoch
                  )
  
file_path: pytorch-optimizer/examples/mnist.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.tensorboard import SummaryWriter
  from torchvision import datasets, transforms, utils
  import torch_optimizer
label_desc: Use image node_id 4 for calling a global function with example usage: main() without return types