GPU and CUDA

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device, torch.cuda.get_device_name(0) if device.type == "cuda" else "cpu")

# Move tensors
x = torch.randn(64, 3).to(device)

# Move model
model = torch.nn.Linear(3, 1).to(device)

# Mixed precision for speed
scaler = torch.cuda.amp.GradScaler(enabled=device.type == "cuda")
with torch.cuda.amp.autocast(enabled=device.type == "cuda"):
    out = model(x)
    loss = out.sum()
scaler.scale(loss).backward()
scaler.step(torch.optim.SGD(model.parameters(), lr=0.01))
scaler.update()

# Multi-GPU
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)

# Clear cache
torch.cuda.empty_cache()