code: def main():
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "--mode", type=str, default="SPEN", choices=["SPEN", "ICNN"]
      )
      args = parser.parse_args()
  
      # Set up the MNIST data loaders.
      train_loader = DataLoader(
          datasets.MNIST(
              "/tmp/mnist-data",
              train=True,
              download=True,
              transform=transforms.Compose(
                  [
                      transforms.ToTensor(),
                      transforms.Normalize((0.1307,), (0.3081,)),
                  ]
              ),
          ),
          batch_size=128,
          shuffle=True,
      )
  
      test_loader = DataLoader(
          datasets.MNIST(
              "/tmp/mnist-data",
              train=False,
              transform=transforms.Compose(
                  [
                      transforms.ToTensor(),
                      transforms.Normalize((0.1307,), (0.3081,)),
                  ]
              ),
          ),
          batch_size=128,
          shuffle=True,
      )
  
      # Initialize an energy model and optimizer for the energy function.
      device = torch.device("cuda")
      Enet = EnergyNet().to(device)
  
      if args.mode == "ICNN":
          # In the ICNN, the fce2 weights are required to be non-negative
          # to maintain convexity of the energy function.
          # Here we just project negative weights back onto the
          # non-negative orthant.
          fce2W = Enet.state_dict()["fce2.weight"]
          fce2W[fce2W < 0] = 0.0
  
      model = UnrollEnergy(Enet).to(device)
      optimizer = optim.Adam(Enet.parameters(), lr=3e-4)
  
      for epoch in range(10):
          train(
              model, device, train_loader, optimizer, args.mode, epoch
          )
          test(model, device, test_loader)
  
imports: import argparse
  import typing
  import torch
  from torch import nn
  from torch.utils.data import DataLoader
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  import higher
label_desc: Use image node_id 1 for calling a global function with example usage: main() without return types