InterviewsVector

Getting started with PyTorch

Python

Quick answer

PyTorch is a deep learning framework built around tensors (like NumPy arrays but GPU-accelerated) and automatic differentiation via autograd. Install it with the official command from pytorch.org for your CUDA version, create tensors with torch.tensor, define models by subclassing nn.Module, and train with an optimizer and a loss in a manual loop. Its define-by-run graph makes it easy to debug.

Short answer: PyTorch is a deep-learning framework built on two ideas: tensors (NumPy-like arrays that run on the GPU) and autograd (automatic differentiation). Install the exact command from pytorch.org for your OS/CUDA, create tensors with torch.tensor, define models by subclassing nn.Module, and train in a manual loop with an optimizer and a loss. Its define-by-run graph makes it easy to debug with plain Python.

Install

Don't guess the install command — PyTorch's selector gives you the exact line for your OS, package manager, and CUDA version. The common pip forms:

# CPU-only
pip install torch torchvision
 
# CUDA 12.x build (example — check the selector for the current index URL)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

Verify it imports and detects your accelerator:

import torch
 
print(torch.__version__)
print(torch.cuda.is_available())     # True if a CUDA GPU is set up

Tensors

Tensors are PyTorch's core data structure — like NumPy arrays, but they can live on a GPU and track gradients. Use lowercase torch.tensor (the factory that infers dtype), not the legacy torch.Tensor constructor:

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
r = torch.rand(3, 3)                 # random values in [0, 1)
z = torch.zeros(2, 4)                # a 2x4 tensor of zeros
 
print(x.shape, x.dtype)              # torch.Size([2, 3]) torch.float32
print(x.mean(dim=1))                 # tensor([2., 5.]) — mean of each row

Move a tensor (or a model) to the GPU with .to(device) — the standard portable pattern:

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

Autograd: the reason PyTorch exists

Set requires_grad=True and PyTorch records operations so it can compute gradients with .backward(). This is what powers training:

w = torch.tensor(3.0, requires_grad=True)
y = w ** 2 + 2 * w          # y = w² + 2w
y.backward()                # compute dy/dw
print(w.grad)               # tensor(8.) — since dy/dw = 2w + 2 = 8 at w=3

You rarely call .backward() by hand on scalars like this — it runs inside the training loop below — but this is the mechanism underneath.

Define a model with nn.Module

Subclass nn.Module, create layers in __init__, and implement forward:

import torch.nn as nn
 
class MLP(nn.Module):
    def __init__(self, in_features, hidden, out_features):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_features),
        )
 
    def forward(self, x):
        return self.net(x)
 
model = MLP(10, 32, 1).to(device)

The training loop

PyTorch doesn't hide the loop the way Keras .fit() does — you write it, which is exactly why it's easy to debug. The four steps are: zero the gradients, forward, backward, step.

import torch.nn as nn
 
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()
 
# dummy data
X = torch.rand(64, 10, device=device)
target = torch.rand(64, 1, device=device)
 
for epoch in range(100):
    optimizer.zero_grad()        # 1. clear old gradients
    pred = model(X)              # 2. forward pass
    loss = loss_fn(pred, target)
    loss.backward()              # 3. autograd computes gradients
    optimizer.step()             # 4. update the weights
 
    if epoch % 20 == 0:
        print(epoch, loss.item())

Forgetting optimizer.zero_grad() is the classic beginner bug: PyTorch accumulates gradients by default, so without it your gradients from every batch add up and training diverges.

Inference

Switch to eval mode and disable gradient tracking for prediction — it's faster and uses less memory:

model.eval()
with torch.no_grad():
    preds = model(X)

Related guides

Sources

Key takeaways

  • •PyTorch is a deep learning framework built on tensors (NumPy-like but GPU-accelerated) and autograd for automatic differentiation.
  • •Install it with the official command from pytorch.org that matches your CUDA version.
  • •Define models by subclassing nn.Module and train in a manual loop with an optimizer and a loss.
  • •Its define-by-run dynamic graph makes debugging straightforward.

Frequently asked questions

What is PyTorch?

A deep learning framework centred on GPU-accelerated tensors and autograd, with dynamic define-by-run graphs.

How do I install PyTorch?

Use the official selector command from pytorch.org for your OS and CUDA version rather than a generic pip install.

How do you define a model in PyTorch?

Subclass nn.Module, create the layers in __init__, and implement the forward pass in forward().

By Mohammad Wasi

Software Engineering Leader & Technical Author · Updated September 9, 2026

See all posts

Related Posts