Skip to content
Deep Learning with PyTorch
Neural Networks 8 min read

PyTorch & Tensors

Create and reshape tensors, bridge to NumPy, move to the GPU, and let autograd compute derivatives for you.

Download notebook Open Google ColabIn Colab: File → Upload notebook → pick the downloaded file.

PyTorch is the workhorse of modern deep learning research and much of industry. At its core it gives you three things: tensors (NumPy-like arrays), GPU acceleration (the same code runs on a graphics card), and autograd (automatic gradients — the machinery that makes training possible). This lesson covers all three.

Run this lesson in Colab

Download this lesson as a notebook and run it in Google Colab (free at colab.research.google.com) — PyTorch is preinstalled there. The small browser demos below still run right here.

Creating tensors

A tensor is an n-dimensional array: a scalar is 0-D, a vector 1-D, a matrix 2-D, and a color image batch is 4-D. Everything in PyTorch is a tensor.

import torch
import numpy as np
 
# From Python data
a = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])
 
# Factory functions
zeros = torch.zeros(2, 3)          # all zeros
ones  = torch.ones(5)              # all ones
noise = torch.randn(3, 4)          # standard normal random values
 
print(a)
print(a.shape)                     # torch.Size([2, 3])
print(a.dtype)                     # torch.float32

Two attributes you'll check constantly:

  • .shape — the size along each dimension. Most PyTorch bugs are shape bugs.
  • .dtype — the element type. float32 is the default for model inputs and weights; class labels must be integers (int64, a.k.a. torch.long).
labels = torch.tensor([0, 2, 1])       # int64 automatically
floats = torch.tensor([0, 2, 1], dtype=torch.float32)
print(labels.dtype, floats.dtype)

Reshaping, matmul, and friends

Networks constantly reshape data — flattening images into vectors, adding batch dimensions, reordering channels:

x = torch.arange(12, dtype=torch.float32)   # shape [12]
 
m = x.view(3, 4)          # reshape to 3x4 (same memory)
m2 = x.reshape(3, -1)     # -1 means "infer this dimension"
print(m.shape, m2.shape)
 
img = x.view(1, 3, 2, 2)              # [batch, channels, height, width]
hwc = img.permute(0, 2, 3, 1)         # reorder dims -> [1, 2, 2, 3]
print(hwc.shape)
 
col = x.unsqueeze(1)      # add a dimension: [12] -> [12, 1]
back = col.squeeze()      # drop size-1 dimensions: [12, 1] -> [12]
print(col.shape, back.shape)

Matrix multiplication uses the @ operator — this is the single most important operation in deep learning (a linear layer is one matmul plus a bias):

W = torch.randn(4, 3)     # weights: 3 inputs -> 4 outputs
x = torch.randn(3)        # one sample with 3 features
b = torch.randn(4)
 
y = W @ x + b             # a linear layer, by hand
print(y.shape)            # torch.Size([4])

Also handy: max vs argmax. Given a batch of class scores, max(1) gives the highest score per row, while argmax(1) gives which class had it — that's how you turn network outputs into predictions:

scores = torch.tensor([[0.1, 0.7, 0.2],
                       [0.8, 0.1, 0.1]])
print(scores.argmax(1))       # tensor([1, 0]) -> predicted classes
values, indices = scores.max(1)
print(values, indices)

The NumPy bridge

Tensors and NumPy arrays convert both ways cheaply (they can even share memory on CPU):

arr = np.array([[1.0, 2.0], [3.0, 4.0]])
 
t = torch.from_numpy(arr)     # numpy -> tensor
back = t.numpy()              # tensor -> numpy
 
single = torch.tensor(3.14)
print(single.item())          # one-element tensor -> Python float

Moving to the GPU

The magic of PyTorch: the same code runs on CPU or GPU. The modern idiom is to define a device once and move both model and data to it:

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
 
x = torch.randn(1000, 1000).to(device)
y = x @ x                     # runs on the GPU if one is available

In Colab, enable the GPU via Runtime → Change runtime type → T4 GPU. One rule to remember: everything in one operation must live on the same device — mixing a CPU tensor with a GPU tensor raises an error. And to plot or convert to NumPy, bring data back with .cpu() first.

Autograd: derivatives for free

Here's the feature that separates PyTorch from NumPy. Mark a tensor with requires_grad=True and PyTorch records every operation done to it. Call .backward() on a scalar result, and it walks the recorded graph in reverse (backpropagation) to fill in .grad — the derivative of the result with respect to each input.

Take y = x² + 3x. Calculus says dy/dx = 2x + 3, so at x = 2 the derivative is 7. Autograd agrees:

x = torch.tensor(2.0, requires_grad=True)
 
y = x**2 + 3*x        # forward pass: PyTorch records the graph
y.backward()          # backward pass: compute dy/dx
 
print(x.grad)         # tensor(7.) -> matches 2*2 + 3

No symbolic math, no hand-derived formulas — and it scales to functions with millions of inputs, which is exactly what a neural network's loss is.

For intuition, here's the same derivative computed numerically the way you might in a first calculus course — nudge x a tiny bit and measure the change. Run it in your browser:

Python — runs in your browser

Numerical differentiation needs one extra function evaluation per parameter — hopeless for a million-parameter model. Autograd gets every gradient in a single backward pass. That efficiency is why deep learning is feasible at all.

One more autograd detail: gradients accumulate. Each .backward() call adds to .grad rather than replacing it. That's why training loops zero the gradients every step — you'll meet optimizer.zero_grad() in the next lesson.

Turning autograd off

Tracking operations costs memory and time. When you're only using a model (predicting, evaluating), wrap the code so no graph is recorded:

x = torch.tensor(2.0, requires_grad=True)
 
with torch.no_grad():                 # classic way
    y = x**2 + 3*x
print(y.requires_grad)                # False
 
with torch.inference_mode():          # modern, slightly faster
    y = x**2 + 3*x
print(y.requires_grad)                # False

Prefer torch.inference_mode() for evaluation and deployment code — it's the stricter, faster successor to no_grad.

Check your understanding

5 questions · free
  1. Q1.A batch of 32 RGB images of size 28x28 is stored as a tensor. What is its conventional shape in PyTorch?

  2. Q2.What does x.view(4, -1) do to a tensor with 12 elements?

  3. Q3.After y = x**3 with x = torch.tensor(2.0, requires_grad=True) and y.backward(), what is x.grad?

  4. Q4.Why do we wrap evaluation code in torch.inference_mode()?

  5. Q5.You get 'Expected all tensors to be on the same device'. What's the likely cause?

Exercise: Autograd vs. calculus

For the function y = 2x³ − 5x² + 4x, compute dy/dx at x = 3 three ways: (1) with autograd, (2) analytically by hand (write the derivative formula in code), and confirm they match. Then demonstrate gradient accumulation: call .backward() a second time (rebuild y first) and print x.grad — explain what you see, then fix it with x.grad.zero_().

Next up: the training loop — model, criterion, optimizer, and the forward-backward-step dance that turns gradients into learning.