GPU Devices

An advance of using PyTorch to numpy is the ability to leverage device level specifications for computation. GPU tend to perform way better than CPUs in training neural networks. Therefore, being able to cast tensors to device type is very beneficial and common practice.

Device Type MPS

Since I use a MacBook with M1/M2 chips, the device specification for my needs is mps. For non-Mac users, this is likely to be CUDA. Use this command instead: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

import torch

# macbook M1-M2 chips specification
device = torch.device('mps') 
x = torch.tensor([1, 2, 3])
y = torch.tensor([1, 2, 3])

# casting tensors x and y to device
x.to(device), y.to(device)
(tensor([1, 2, 3], device='mps:0'), tensor([1, 2, 3], device='mps:0'))

Implementing on device operation

x.add(y)
tensor([2, 4, 6])