Add Apple M-series (MPS) device support to Python scripts

[brianholland]

Summary

Most standalone Python scripts (.py) only detect CUDA or fall back to CPU, ignoring Apple Silicon's MPS backend. This means users with M1/M2/M3/M4 Macs always fall back to CPU, missing significant GPU acceleration.

Problem

~29 Python files use this pattern:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

This should be updated to also detect MPS:

if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

Additional considerations

• pin_memory=True in DataLoaders should be conditional on CUDA availability
• torch.compile should be guarded (limited MPS support)
• DDP scripts are inherently CUDA-only and should be left as-is

Scope

• All standalone .py training/inference scripts in ch04-ch07, appendix, and pkg/tests
• Excludes DDP-specific scripts (appendix-A DDP scripts, ch05/10 multi-GPU)
• Excludes notebooks (many already have MPS support)

🤖 Generated with Claude Code

[rasbt]

If someone reads this, this omission of mps was on purpose (see #625 #644 #442 #846)

Perhaps worthwhile revisiting some time in the future though.

[amit-chaubey]

if you use mps

torch.manual_seed(123)
model = NeuralNetwork(num_inputs=2, num_outputs=2)

device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
model.to(device)

optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

num_epochs = 3
start_time = time.time()
for epoch in range(num_epochs):

model.train()
for batch_idx, (features, labels) in enumerate(train_loader):
    features, labels = features.to(device), labels.to(device)
    logits = model(features)
    loss = F.cross_entropy(logits, labels) # Loss function

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    ### LOGGING
    print(f"Epoch: {epoch+1:03d}/{num_epochs:03d}"
          f" | Batch {batch_idx:03d}/{len(train_loader):03d}"
          f" | Train/Val Loss: {loss:.2f}")


model.eval()

end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")
#Time taken: 0.014540672302246094 seconds

---

if you use cpu

torch.manual_seed(123)
model = NeuralNetwork(num_inputs=2, num_outputs=2)

device = torch.device("cpu")
model.to(device)

optimizer = torch.optim.SGD(model.parameters(), lr=0.5)

num_epochs = 3

start_time = time.time()
for epoch in range(num_epochs):

model.train()
for batch_idx, (features, labels) in enumerate(train_loader):
    features, labels = features.to(device), labels.to(device)
    logits = model(features)
    loss = F.cross_entropy(logits, labels) # Loss function

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    ### LOGGING
    print(f"Epoch: {epoch+1:03d}/{num_epochs:03d}"
          f" | Batch {batch_idx:03d}/{len(train_loader):03d}"
          f" | Train/Val Loss: {loss:.2f}")
model.eval()

end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")

#Time taken: 0.0030498504638671875 seconds

MPS is stuck in no-man's land, too slow for large models, overkill overhead for small ones. i feel its much slower.

you may use runpod and any other cloud gpu to accelerate training. one of the reason could be Apple unified memory is shared between CPU/GPU.

Thanks for amazing book @rasbt