PyTorch 2 cheatsheet

PyTorch 2: Tensors

8 common PyTorch 2 tasks.

Install PyTorch

Install the library. Choose your CUDA version at pytorch.org/get-started for GPU builds.

pip install torch torchvision

The standard imports for building models.

import torch
import torch.nn as nn

Create tensors from data, zeros, ones or random values.

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
zeros = torch.zeros(3, 4)
rand = torch.rand(2, 3)

Inspect a tensor's properties.

x.shape, x.dtype, x.device

Change the shape. -1 infers one dimension.

x.reshape(-1, 2)
x.unsqueeze(0)
x.squeeze()

Share memory with NumPy arrays on the CPU.

t = torch.from_numpy(arr)
arr = t.numpy()

Matrix multiplication with the @ operator.

y = a @ b

Use a GPU if available, including Apple Silicon.

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