Update patch for MNIST main.py, as main.py has been updated

git hash of the latest main.py is d0b7e37

Author: wkliao

examples/MNIST/mnist.patch

@@ -1,5 +1,5 @@
---- mnist_main.py.orig	2025-05-09 10:51:06.814200110 -0500
-+++ mnist_main.py	2025-05-09 11:15:17.198167820 -0500
+--- main.py.d0b7e37	2026-02-18 14:50:27.722389236 -0600
++++ main.py	2026-02-18 14:59:02.737245645 -0600
 @@ -1,3 +1,8 @@
 +#
 +# Copyright (C) 2025, Northwestern University and Argonne National Laboratory
@@ -55,7 +55,7 @@
      args = parser.parse_args()
 
      use_accel = not args.no_accel and torch.accelerator.is_available()
-@@ -103,12 +119,11 @@
+@@ -103,7 +119,7 @@
      else:
          device = torch.device("cpu")
 
@@ -64,13 +64,7 @@
      test_kwargs = {'batch_size': args.test_batch_size}
      if use_accel:
          accel_kwargs = {'num_workers': 1,
--                       'pin_memory': True,
--                       'shuffle': True}
-+                       'pin_memory': True}
-         train_kwargs.update(accel_kwargs)
-         test_kwargs.update(accel_kwargs)
- 
-@@ -116,25 +131,53 @@
+@@ -117,25 +133,54 @@
          transforms.ToTensor(),
          transforms.Normalize((0.1307,), (0.3081,))
          ])
@@ -115,19 +109,21 @@
 -        torch.save(model.state_dict(), "mnist_cnn.pt")
 +        if rank == 0:
 +            torch.save(model.state_dict(), "mnist_cnn.pt")
- 
++
 +    # close files
 +    train_file.close()
 +    test_file.close()
 
+ 
  if __name__ == '__main__':
+-    main()
 +    ## initialize parallel environment
 +    comm, device = comm_file.init_parallel()
 +
 +    rank = comm.get_rank()
 +    nprocs = comm.get_size()
 +
-     main()
++    main()
 +
 +    comm.finalize()
 +