MPS support updates

tasks/image_classification/train.py

@@ -212,8 +212,13 @@ if __name__=='__main__':
     with open(f'{args.log_dir}/args.txt', 'w') as f:
         print(args, file=f)
 
-    # Configure device string
-    device = f'cuda:{args.device[0]}' if args.device[0] != -1 else 'cpu'
+    # Configure device string (support MPS on macOS)
+    if args.device[0] != -1:
+        device = f'cuda:{args.device[0]}'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        device = 'cpu'
     print(f'Running model {args.model} on {device}')
 
     # Build model conditionally

tasks/mazes/train.py

@@ -151,9 +151,15 @@ if __name__=='__main__':
     with open(f'{args.log_dir}/args.txt', 'w') as f:
         print(args, file=f)
 
-    # Configure device string
-    device = f'cuda:{args.device[0]}' if args.device[0] != -1 else 'cpu'
-    print(f'Running model {args.model} on {device} for dataset {args.dataset}')
+    # Configure device string (support MPS on macOS)
+    if args.device[0] != -1:
+        device = f'cuda:{args.device[0]}'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        device = 'cpu'
+    print(f'Running model {args.model} on {device}')
+
 
     # Build model conditionally
     model = None

tasks/parity/train.py

@@ -112,8 +112,14 @@ if __name__=='__main__':
     with open(f'{args.log_dir}/args.txt', 'w') as f:
         print(args, file=f)  
 
-    # Configure device string
-    device = f'cuda:{args.device[0]}' if args.device[0] != -1 else 'cpu' 
+    # Configure device string (support MPS on macOS)
+    if args.device[0] != -1:
+        device = f'cuda:{args.device[0]}'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        device = 'cpu'
+    print(f'Running model {args.model} on {device}')
 
     # Build model
     model = prepare_model(prediction_reshaper, args, device)

tasks/qamnist/train.py

@@ -141,9 +141,14 @@ if __name__=='__main__':
     with open(f'{args.log_dir}/args.txt', 'w') as f:
         print(args, file=f)  
 
-    # Configure device string
-    device = f'cuda:{args.device[0]}' if args.device[0] != -1 else 'cpu' 
-    print(f'Running on {device}')
+    # Configure device string (support MPS on macOS)
+    if args.device[0] != -1:
+        device = f'cuda:{args.device[0]}'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        device = 'cpu'
+    print(f'Running model {args.model} on {device}')
 
     # Build model
     model = prepare_model(args, device)

tasks/sort/train.py

@@ -154,9 +154,14 @@ if __name__=='__main__':
     with open(f'{args.log_dir}/args.txt', 'w') as f:
         print(args, file=f)  
 
-    # Configure device string
-    device = f'cuda:{args.device[0]}' if args.device[0] != -1 else 'cpu' 
-    print(f'Running on {device}')
+    # Configure device string (support MPS on macOS)
+    if args.device[0] != -1:
+        device = f'cuda:{args.device[0]}'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        device = 'cpu'
+    print(f'Running model {args.model} on {device}')
 
     # Build model
     model = ContinuousThoughtMachineSORT(