prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  Add `save_model` arg to `mnist_hogwild` example
  Currently the example doesn't support the `--save_model` argument like the other examples
  
  
  </issue>
  <code>
  [start of mnist_hogwild/main.py]
  1 from __future__ import print_function
  2 import argparse
  3 import torch
  4 import torch.nn as nn
  5 import torch.nn.functional as F
  6 import torch.multiprocessing as mp
  7 from torch.utils.data.sampler import Sampler
  8 from torchvision import datasets, transforms
  9 
  10 from train import train, test
  11 
  12 # Training settings
  13 parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
  14 parser.add_argument('--batch-size', type=int, default=64, metavar='N',
  15                     help='input batch size for training (default: 64)')
  16 parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
  17                     help='input batch size for testing (default: 1000)')
  18 parser.add_argument('--epochs', type=int, default=10, metavar='N',
  19                     help='number of epochs to train (default: 10)')
  20 parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
  21                     help='learning rate (default: 0.01)')
  22 parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
  23                     help='SGD momentum (default: 0.5)')
  24 parser.add_argument('--seed', type=int, default=1, metavar='S',
  25                     help='random seed (default: 1)')
  26 parser.add_argument('--log-interval', type=int, default=10, metavar='N',
  27                     help='how many batches to wait before logging training status')
  28 parser.add_argument('--num-processes', type=int, default=2, metavar='N',
  29                     help='how many training processes to use (default: 2)')
  30 parser.add_argument('--cuda', action='store_true', default=False,
  31                     help='enables CUDA training')
  32 parser.add_argument('--mps', action='store_true', default=False,
  33                         help='enables macOS GPU training')
  34 parser.add_argument('--dry-run', action='store_true', default=False,
  35                     help='quickly check a single pass')
  36 
  37 class Net(nn.Module):
  38     def __init__(self):
  39         super(Net, self).__init__()
  40         self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
  41         self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
  42         self.conv2_drop = nn.Dropout2d()
  43         self.fc1 = nn.Linear(320, 50)
  44         self.fc2 = nn.Linear(50, 10)
  45 
  46     def forward(self, x):
  47         x = F.relu(F.max_pool2d(self.conv1(x), 2))
  48         x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
  49         x = x.view(-1, 320)
  50         x = F.relu(self.fc1(x))
  51         x = F.dropout(x, training=self.training)
  52         x = self.fc2(x)
  53         return F.log_softmax(x, dim=1)
  54 
  55 
  56 if __name__ == '__main__':
  57     args = parser.parse_args()
  58 
  59     use_cuda = args.cuda and torch.cuda.is_available()
  60     use_mps = args.mps and torch.backends.mps.is_available()
  61     if use_cuda:
  62         device = torch.device("cuda")
  63     elif use_mps:
  64         device = torch.device("mps")
  65     else:
  66         device = torch.device("cpu")
  67 
  68     transform=transforms.Compose([
  69         transforms.ToTensor(),
  70         transforms.Normalize((0.1307,), (0.3081,))
  71         ])
  72     dataset1 = datasets.MNIST('../data', train=True, download=True,
  73                        transform=transform)
  74     dataset2 = datasets.MNIST('../data', train=False,
  75                        transform=transform)
  76     kwargs = {'batch_size': args.batch_size,
  77               'shuffle': True}
  78     if use_cuda:
  79         kwargs.update({'num_workers': 1,
  80                        'pin_memory': True,
  81                       })
  82 
  83     torch.manual_seed(args.seed)
  84     mp.set_start_method('spawn', force=True)
  85 
  86     model = Net().to(device)
  87     model.share_memory() # gradients are allocated lazily, so they are not shared here
  88 
  89     processes = []
  90     for rank in range(args.num_processes):
  91         p = mp.Process(target=train, args=(rank, args, model, device,
  92                                            dataset1, kwargs))
  93         # We first train the model across `num_processes` processes
  94         p.start()
  95         processes.append(p)
  96     for p in processes:
  97         p.join()
  98 
  99     # Once training is complete, we can test the model
  100     test(args, model, device, dataset2, kwargs)
  101
  [end of mnist_hogwild/main.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/mnist_hogwild/main.py b/mnist_hogwild/main.py
  --- a/mnist_hogwild/main.py
  +++ b/mnist_hogwild/main.py
  @@ -30,7 +30,9 @@
   parser.add_argument('--cuda', action='store_true', default=False,
                       help='enables CUDA training')
   parser.add_argument('--mps', action='store_true', default=False,
  -                        help='enables macOS GPU training')
  +                    help='enables macOS GPU training')
  +parser.add_argument('--save_model', action='store_true', default=False,
  +                    help='save the trained model to state_dict')
   parser.add_argument('--dry-run', action='store_true', default=False,
                       help='quickly check a single pass')
   
  @@ -96,5 +98,8 @@
       for p in processes:
           p.join()
   
  +    if args.save_model:
  +        torch.save(model.state_dict(), "MNIST_hogwild.pt")
  +
       # Once training is complete, we can test the model
       test(args, model, device, dataset2, kwargs)
  
verification_info: {"golden_diff": "diff --git a/mnist_hogwild/main.py b/mnist_hogwild/main.py\n--- a/mnist_hogwild/main.py\n+++ b/mnist_hogwild/main.py\n@@ -30,7 +30,9 @@\n parser.add_argument('--cuda', action='store_true', default=False,\n                     help='enables CUDA training')\n parser.add_argument('--mps', action='store_true', default=False,\n-                        help='enables macOS GPU training')\n+                    help='enables macOS GPU training')\n+parser.add_argument('--save_model', action='store_true', default=False,\n+                    help='save the trained model to state_dict')\n parser.add_argument('--dry-run', action='store_true', default=False,\n                     help='quickly check a single pass')\n \n@@ -96,5 +98,8 @@\n     for p in processes:\n         p.join()\n \n+    if args.save_model:\n+        torch.save(model.state_dict(), \"MNIST_hogwild.pt\")\n+\n     # Once training is complete, we can test the model\n     test(args, model, device, dataset2, kwargs)\n", "issue": "Add `save_model` arg to `mnist_hogwild` example\nCurrently the example doesn't support the `--save_model` argument like the other examples\r\n\n", "before_files": [{"content": "from __future__ import print_function\nimport argparse\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.multiprocessing as mp\nfrom torch.utils.data.sampler import Sampler\nfrom torchvision import datasets, transforms\n\nfrom train import train, test\n\n# Training settings\nparser = argparse.ArgumentParser(description='PyTorch MNIST Example')\nparser.add_argument('--batch-size', type=int, default=64, metavar='N',\n                    help='input batch size for training (default: 64)')\nparser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n                    help='input batch size for testing (default: 1000)')\nparser.add_argument('--epochs', type=int, default=10, metavar='N',\n                    help='number of epochs to train (default: 10)')\nparser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n                    help='learning rate (default: 0.01)')\nparser.add_argument('--momentum', type=float, default=0.5, metavar='M',\n                    help='SGD momentum (default: 0.5)')\nparser.add_argument('--seed', type=int, default=1, metavar='S',\n                    help='random seed (default: 1)')\nparser.add_argument('--log-interval', type=int, default=10, metavar='N',\n                    help='how many batches to wait before logging training status')\nparser.add_argument('--num-processes', type=int, default=2, metavar='N',\n                    help='how many training processes to use (default: 2)')\nparser.add_argument('--cuda', action='store_true', default=False,\n                    help='enables CUDA training')\nparser.add_argument('--mps', action='store_true', default=False,\n                        help='enables macOS GPU training')\nparser.add_argument('--dry-run', action='store_true', default=False,\n                    help='quickly check a single pass')\n\nclass Net(nn.Module):\n    def __init__(self):\n        super(Net, self).__init__()\n        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n        self.conv2_drop = nn.Dropout2d()\n        self.fc1 = nn.Linear(320, 50)\n        self.fc2 = nn.Linear(50, 10)\n\n    def forward(self, x):\n        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n        x = x.view(-1, 320)\n        x = F.relu(self.fc1(x))\n        x = F.dropout(x, training=self.training)\n        x = self.fc2(x)\n        return F.log_softmax(x, dim=1)\n\n\nif __name__ == '__main__':\n    args = parser.parse_args()\n\n    use_cuda = args.cuda and torch.cuda.is_available()\n    use_mps = args.mps and torch.backends.mps.is_available()\n    if use_cuda:\n        device = torch.device(\"cuda\")\n    elif use_mps:\n        device = torch.device(\"mps\")\n    else:\n        device = torch.device(\"cpu\")\n\n    transform=transforms.Compose([\n        transforms.ToTensor(),\n        transforms.Normalize((0.1307,), (0.3081,))\n        ])\n    dataset1 = datasets.MNIST('../data', train=True, download=True,\n                       transform=transform)\n    dataset2 = datasets.MNIST('../data', train=False,\n                       transform=transform)\n    kwargs = {'batch_size': args.batch_size,\n              'shuffle': True}\n    if use_cuda:\n        kwargs.update({'num_workers': 1,\n                       'pin_memory': True,\n                      })\n\n    torch.manual_seed(args.seed)\n    mp.set_start_method('spawn', force=True)\n\n    model = Net().to(device)\n    model.share_memory() # gradients are allocated lazily, so they are not shared here\n\n    processes = []\n    for rank in range(args.num_processes):\n        p = mp.Process(target=train, args=(rank, args, model, device,\n                                           dataset1, kwargs))\n        # We first train the model across `num_processes` processes\n        p.start()\n        processes.append(p)\n    for p in processes:\n        p.join()\n\n    # Once training is complete, we can test the model\n    test(args, model, device, dataset2, kwargs)\n", "path": "mnist_hogwild/main.py"}]}