ViteBetti with torch can use Apple's AMD GPUs

Author: sambit-giri

src/tools21cm/ViteBetti.py

@@ -34,7 +34,12 @@ def jit(func, *args, **kwargs):
 try:
     import torch
     torch_available = True
-    torch_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if torch.cuda.is_available():
+        torch_device = torch.device('cuda')
+    elif torch.backends.mps.is_available():
+        torch_device = torch.device('mps')
+    else:
+        torch_device = torch.device('cpu')
 except ImportError:
     torch_available = False
     torch = None

src/tools21cm/topology.py

@@ -21,7 +21,7 @@ def EulerCharacteristic(data, thres=0.5, neighbors=6, speed_up='cython', verbose
 	speed_up: str
 		Method used to speed up calculation (Default: cython). 
 		Options are cython, numba, torch and numpy.
-		The caclulation with torch uses GPUs if available on the device.
+		The calculation with torch uses GPUs if available on the device.
 	verbose: bool
 		If True, verbose is printed.