844fa63c841e6bed6d05606445640486a6ca4dd0,custom_containers/pytorch/mnist/trainer/mnist.py,,main,#,93

Before Change

def main():
  // Training settings
  parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
  parser.add_argument(
      "--batch-size",
      type=int,
      default=64,
      metavar="N",
      help="input batch size for training (default: 64)")
  parser.add_argument(
      "--test-batch-size",
      type=int,
      default=1000,
      metavar="N",
      help="input batch size for testing (default: 1000)")
  parser.add_argument(
      "--epochs",
      type=int,
      default=10,
      metavar="N",
      help="number of epochs to train (default: 10)")
  parser.add_argument(
      "--lr",
      type=float,
      default=0.01,
      metavar="LR",
      help="learning rate (default: 0.01)")
  parser.add_argument(
      "--momentum",
      type=float,
      default=0.5,
      metavar="M",
      help="SGD momentum (default: 0.5)")
  parser.add_argument(
      "--no-cuda",
      action="store_true",
      default=False,
      help="disables CUDA training")
  parser.add_argument(
      "--seed",
      type=int,
      default=1,
      metavar="S",
      help="random seed (default: 1)")
  parser.add_argument(
      "--log-interval",
      type=int,
      default=10,
      metavar="N",
      help="how many batches to wait before logging training status")
  parser.add_argument(
      "--model-dir",
      default=None,
      help="The directory to store the model")

  args = parser.parse_args()
  use_cuda = not args.no_cuda and torch.cuda.is_available()

  torch.manual_seed(args.seed)

After Change

def main():
  // Training settings
  args = get_args()
  use_cuda = not args.no_cuda and torch.cuda.is_available()

  torch.manual_seed(args.seed)

Instances