very slight refactor, bit cleaner

train.py

@@ -112,12 +112,10 @@ def get_batch(split):
     ix = torch.randint(len(data) - block_size, (batch_size,))
     x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
     y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
-    # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
-    if "cuda" in device:
-        # GPU training
+    if device_type == 'cuda':
+        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
         x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
     else:
-        # CPU or MPS training
         x, y = x.to(device), y.to(device)
     return x, y