Tensors & autograd in PyTorch

Everything in PyTorch — images, sentences, model weights, gradients — is a tensor: an n-dimensional array, much like a NumPy array but able to live on a GPU and to track the operations performed on it. Mastering tensors and autograd (automatic differentiation) is the foundation for every later lesson.

Creating tensors

import torch

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])      # from a Python list
z = torch.zeros(3, 4)                            # 3x4 of zeros
r = torch.rand(2, 3)                             # uniform [0, 1)
n = torch.randn(2, 3)                            # standard normal
s = torch.arange(0, 10, 2)                       # tensor([0, 2, 4, 6, 8])

print(a.shape, a.dtype, a.device)                # torch.Size([2, 2]) torch.float32 cpu

Ranks and shapes you will meet

Data Typical shape Meaning
Scalar loss () a single number
Daily sales for 30 days (30,) vector
Batch of tabular rows (64, 12) 64 rows, 12 features
Batch of RGB images (32, 3, 224, 224) N, channels, height, width
Batch of token IDs (16, 128) 16 sentences, 128 tokens each

Most bugs in deep learning code are shape bugs, so print .shape often.

Everyday operations

x = torch.randn(64, 12)
w = torch.randn(12, 1)

y = x @ w                 # matrix multiply -> (64, 1)
y = y.squeeze(1)          # drop size-1 dim -> (64,)
imgs = torch.rand(8, 3, 32, 32)
flat = imgs.view(8, -1)   # (8, 3072); -1 lets PyTorch infer the size
mean_per_channel = imgs.mean(dim=(0, 2, 3))   # (3,)
prices = torch.tensor([120.0, 450.0, 80.0])
with_gst = prices * 1.05  # broadcasting a scalar

Broadcasting stretches dimensions of size 1 so shapes match, exactly as in NumPy. Use .reshape() when a tensor may not be contiguous in memory; .view() requires contiguity.

NumPy interop and devices

import numpy as np

arr = np.array([1.0, 2.0, 3.0])
t = torch.from_numpy(arr)        # shares memory with arr
back = t.numpy()

device = "cuda" if torch.cuda.is_available() else "cpu"
t_gpu = t.to(device)
result = (t_gpu * 2).cpu()       # bring back to CPU before .numpy()

Both operands of an operation must be on the same device, otherwise PyTorch raises an error.

Autograd: gradients for free

Training a network means nudging each weight in the direction that reduces the loss. That direction is the gradient. Autograd records every operation on tensors with requires_grad=True in a computation graph, then walks it backwards to compute derivatives.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x          # y = x^3 + 2x
y.backward()                # dy/dx = 3x^2 + 2
print(x.grad)               # tensor(14.)

A tiny linear regression shows the full idea — predicting a sweet shop's daily sales (in thousand INR) from footfall:

torch.manual_seed(0)
footfall = torch.rand(100, 1) * 10
sales = 3.0 * footfall + 5.0 + torch.randn(100, 1) * 0.5

w = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.01

for step in range(2000):
    pred = footfall * w + b
    loss = ((pred - sales) ** 2).mean()
    loss.backward()
    with torch.no_grad():       # update without recording
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad.zero_()              # gradients accumulate otherwise
    b.grad.zero_()

print(w.item(), b.item())       # close to 3.0 and 5.0

Later, nn.Module and torch.optim will hide this boilerplate, but it is exactly what happens underneath.

Common mistakes

  • Forgetting to zero gradients, so they accumulate across steps.
  • Updating weights outside torch.no_grad(), which makes autograd track the update itself.
  • Calling .numpy() on a GPU tensor or on a tensor that requires grad (use .detach().cpu().numpy()).
  • Mixing float64 NumPy data with float32 model weights — convert with .float().

Practice

Create a tensor t = torch.arange(24) and reshape it to (2, 3, 4). Compute the sum over the last dimension and the mean over the first. Then use autograd to find the derivative of f(x) = 5x^2 - 3x + 7 at x = 1.5 and confirm it by hand (answer: 12). Finally, modify the regression loop to use a learning rate of 0.1 and observe what happens to the loss.

B

Written by Bablu Kumar Jha

Founder, BeeProton AI · Mumbai, Maharashtra, India