code: def main(args):
      if args.gpu < 0 or not torch.cuda.is_available():
          device = torch.device("cpu")
      else:
          device = torch.device("cuda", args.gpu)
  
      model = SudokuNN(num_steps=args.steps, edge_drop=args.edge_drop)
  
      if args.do_train:
          if not os.path.exists(args.output_dir):
              os.mkdir(args.output_dir)
          model.to(device)
          train_dataloader = sudoku_dataloader(
              args.batch_size, segment="train"
          )
          dev_dataloader = sudoku_dataloader(
              args.batch_size, segment="valid"
          )
  
          opt = Adam(
              model.parameters(),
              lr=args.lr,
              weight_decay=args.weight_decay,
          )
  
          best_dev_acc = 0.0
          for epoch in range(args.epochs):
              model.train()
              for i, g in enumerate(train_dataloader):
                  g = g.to(device)
                  _, loss = model(g)
                  opt.zero_grad()
                  loss.backward()
                  opt.step()
                  if i % 100 == 0:
                      print(
                          f"Epoch {epoch}, batch {i}, loss {loss.cpu().data}"
                      )
  
              # dev
              print("\n=========Dev step========")
              model.eval()
              dev_loss = []
              dev_res = []
              for g in dev_dataloader:
                  g = g.to(device)
                  target = g.ndata["a"]
                  target = target.view([-1, 81])
  
                  with torch.no_grad():
                      preds, loss = model(g, is_training=False)
                      preds = preds.view([-1, 81])
  
                      for i in range(preds.size(0)):
                          dev_res.append(
                              int(
                                  torch.equal(preds[i, :], target[i, :])
                              )
                          )
  
                      dev_loss.append(loss.cpu().detach().data)
  
              dev_acc = sum(dev_res) / len(dev_res)
              print(f"Dev loss {np.mean(dev_loss)}, accuracy {dev_acc}")
              if dev_acc >= best_dev_acc:
                  torch.save(
                      model.state_dict(),
                      os.path.join(args.output_dir, "model_best.bin"),
                  )
                  best_dev_acc = dev_acc
              print(f"Best dev accuracy {best_dev_acc}\n")
  
          torch.save(
              model.state_dict(),
              os.path.join(args.output_dir, "model_final.bin"),
          )
  
      if args.do_eval:
          model_path = os.path.join(args.output_dir, "model_best.bin")
          if not os.path.exists(model_path):
              raise FileNotFoundError("Saved model not Found!")
  
          model.load_state_dict(torch.load(model_path))
          model.to(device)
  
          test_dataloader = sudoku_dataloader(
              args.batch_size, segment="test"
          )
  
          print("\n=========Test step========")
          model.eval()
          test_loss = []
          test_res = []
          for g in test_dataloader:
              g = g.to(device)
              target = g.ndata["a"]
              target = target.view([-1, 81])
  
              with torch.no_grad():
                  preds, loss = model(g, is_training=False)
                  preds = preds
                  preds = preds.view([-1, 81])
  
                  for i in range(preds.size(0)):
                      test_res.append(
                          int(torch.equal(preds[i, :], target[i, :]))
                      )
  
                  test_loss.append(loss.cpu().detach().data)
  
          test_acc = sum(test_res) / len(test_res)
          print(
              f"Test loss {np.mean(test_loss)}, accuracy {test_acc}\n"
          )
  
file_path: dgl/examples/pytorch/rrn/train_sudoku.py
imports: import argparse
  import os
  import numpy
  import torch
  from sudoku import SudokuNN
  from sudoku_data import sudoku_dataloader
  from torch.optim import Adam
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) without return types