with accelerator API

Author: sumantro93

swin_transformer/README.md

@@ -25,7 +25,7 @@ It includes:
 Install dependencies:
 
 ```bash
-pip install torch torchvision
+pip install -r requirements.txt
 ```
 
 ---
@@ -43,7 +43,7 @@ python swin_transformer.py --epochs 10 --batch-size 64 --lr 0.001
 Testing is done automatically after each epoch. To only test, run with:
 
 ```bash
-python swin_transformer.py --epochs 0
+python swin_transformer.py --epochs 1
 ```
 
 ### Save the model

swin_transformer/swin_transformer.py

@@ -172,9 +172,10 @@ def main():
     parser.add_argument('--save-model', action='store_true')
     args = parser.parse_args()
 
-    use_cuda = torch.cuda.is_available()
-    device = torch.device("cuda" if use_cuda else "cpu")
-
+    use_accel = torch.accelerator.is_available()
+    device = torch.accelerator.current_accelerator() if use_accel else torch.device("cpu")
+    print(f"Using device: {device}")
+    
     torch.manual_seed(args.seed)
 
     transform = transforms.Compose([