PyTorch as a Problem-Solving Tool

Need a small PyTorch prototype before the model gets large?

I agree to the Privacy Policy Shape the Experiment

I do not think PyTorch should be introduced as a neural network library first. That is what it becomes later, after the basic mental model is already in place.

At the beginning, PyTorch is more useful as a way to ask a smaller question: can I turn this problem into numbers, compare a prediction with reality, and make the error smaller?

Can I turn this problem into numbers?
Can I compare a prediction with reality?
Can I make the error smaller?

That is already most of machine learning, at least in the useful version. Not the impressive version with diagrams and architecture names, but the practical version where a program improves because the shape of being wrong has been made explicit.

The problem in this article is deliberately small. We have a few examples of project work, and we want to estimate the cost from the number of hours worked.

hours worked -> expected project cost

The relationship is simple enough to write down:

cost = hourly_rate * hours + fixed_admin_cost

But we will pretend we do not know the hourly rate or the fixed cost. We only have observations, and we want the program to recover the rule from those examples.

2 hours  -> 130
4 hours  -> 210
6 hours  -> 290
8 hours  -> 370

This is not interesting because the math is hard. It is interesting because it gives us a clean path through PyTorch without hiding the important parts behind a model that is too large to inspect.

define the state
write the invariant
represent data as tensors
make a prediction
measure error
let gradients improve the parameters
check the result against the original problem

The model comes after the problem. That order matters, because otherwise PyTorch becomes a place to perform machine learning rituals instead of a tool for making a small numerical contract executable.

Hero image: fluid forms in motion, a metaphor for a loss landscape changing under gradients. 3D render by Logan Voss on Unsplash.

Start with the problem, not the framework

Before importing torch, I want the shape of the program in plain English. The system has examples, each example has an input and an expected output, and the model has two numbers it can change: an hourly rate and a fixed fee.

For every input, the model predicts a cost. The loss tells us how wrong those predictions are, and training should reduce that loss by changing the parameters rather than changing the evidence.

data
parameters
prediction
loss
gradient
update
repeat

There is no neural network yet, no GPU, and no architecture diagram. There is only a small numerical problem with one boundary that I want to keep visible.

Training is allowed to change parameters, not the data or the target.

This sounds too basic to write down, but a lot of machine learning bugs are violations of that idea in disguise. The labels leak into the input, the evaluation set becomes part of tuning, a normalization step is fit on the wrong data, or a benchmark quietly becomes the training objective.

At this scale, we can keep the boundary visible. The point of the toy problem is not to make PyTorch look small; it is to make the contract around the code hard to miss.

Docs · PyTorch Tutorials

Automatic Differentiation with torch.autograd

It supports automatic computation of gradient for any computational graph.

— PyTorch Tutorials

Read the PyTorch tutorial

Source card image: a deliberately strange visual for the moment a computation becomes perceptible. 3D render by Milad Fakurian on Unsplash.

Tests as a way to think

I still like writing tests before the mechanism, not because every small notebook needs a full test suite, but because tests force the shape of the problem to become explicit. They make the program say what it is trying to preserve.

For this exercise, I want four properties. The first one is almost embarrassingly small, which is why it is useful: the prediction formula should do what it says.

def test_linear_prediction() -> None:
    hours = torch.tensor([2.0, 4.0])
    rate = torch.tensor(40.0)
    fixed = torch.tensor(50.0)

    prediction = predict_cost(hours, rate, fixed)

    assert torch.allclose(prediction, torch.tensor([130.0, 210.0]))

The second test gives the loss function a clean edge case. If the predictions are perfect, the loss should be zero, because there is no error left to explain.

def test_zero_loss_for_perfect_prediction() -> None:
    prediction = torch.tensor([130.0, 210.0])
    target = torch.tensor([130.0, 210.0])

    assert mean_squared_error(prediction, target).item() == 0.0

The third test checks the learning behavior instead of only checking the formula. Training should move the parameters toward the rule that generated the data.

def test_training_recovers_simple_rule() -> None:
    hours = torch.tensor([2.0, 4.0, 6.0, 8.0])
    costs = torch.tensor([130.0, 210.0, 290.0, 370.0])

    rate, fixed = train_linear_cost_model(hours, costs, steps=1_500, lr=0.01)

    assert abs(rate.item() - 40.0) < 0.5
    assert abs(fixed.item() - 50.0) < 1.0

The fourth test asks for a new prediction. That matters because the goal is not to memorize the four examples; the goal is to recover a rule that still makes sense when the input changes.

def test_trained_model_predicts_new_cost() -> None:
    hours = torch.tensor([2.0, 4.0, 6.0, 8.0])
    costs = torch.tensor([130.0, 210.0, 290.0, 370.0])
    rate, fixed = train_linear_cost_model(hours, costs, steps=1_500, lr=0.01)

    prediction = predict_cost(torch.tensor([10.0]), rate, fixed)

    assert abs(prediction.item() - 450.0) < 2.0

Those tests already describe most of the program. PyTorch is only the tool we use to make the description executable.

Tensors are the boring part, which is good

A tensor is a container for numbers. That definition is incomplete, but it is the useful starting point here, because this exercise does not need anything more dramatic at first.

For this problem, the inputs are a one-dimensional tensor:

hours = torch.tensor([2.0, 4.0, 6.0, 8.0])

The targets are another one-dimensional tensor:

costs = torch.tensor([130.0, 210.0, 290.0, 370.0])

Each position matches the same example. The first input belongs to the first target, the second input belongs to the second target, and so on.

hours[0] -> costs[0]
hours[1] -> costs[1]
hours[2] -> costs[2]
hours[3] -> costs[3]

This is the first real habit to build with PyTorch: know the shape of your data. Not approximately, and not in the way you hope it looks, but exactly.

assert hours.shape == torch.Size([4])
assert costs.shape == torch.Size([4])

When the shape is wrong, the model may still run, and that is the dangerous part. PyTorch will often broadcast tensors for you, which is useful when intended and confusing when accidental.

So before thinking about gradients, I want the tensors to say what they mean. The shape is not decoration around the data; it is part of the data’s meaning.

Tensors become useful when their shape gives the numbers a passage into a new prediction. 3D render by Azza Al Ghardaqa on Unsplash.

The smallest model

The prediction function is just the formula. Keeping it this plain is useful because we can separate the problem from the machinery that will later update the parameters.

import torch


def predict_cost(
    hours: torch.Tensor,
    hourly_rate: torch.Tensor,
    fixed_fee: torch.Tensor,
) -> torch.Tensor:
    return hourly_rate * hours + fixed_fee

If the hourly rate is 40 and the fixed fee is 50, the answer is perfect. At this point we are not learning anything yet; we are only proving that the model can express the rule we care about.

hours = torch.tensor([2.0, 4.0, 6.0, 8.0])
costs = torch.tensor([130.0, 210.0, 290.0, 370.0])

prediction = predict_cost(
    hours,
    hourly_rate=torch.tensor(40.0),
    fixed_fee=torch.tensor(50.0),
)

print(prediction)

Output:

tensor([130., 210., 290., 370.])

So far, PyTorch is not doing anything magical, and that is the point. A model is a function with parameters, and for this problem the parameters are only two numbers.

hourly_rate
fixed_fee

The learning problem is also simple to state:

Find parameter values that make prediction close to target.

That sentence is worth keeping close. When the model gets larger, the same sentence still describes the job.

Loss is the shape of being wrong

The model needs a way to know that one set of parameters is better than another. That is what the loss function is: a numerical definition of being wrong.

For this exercise, mean squared error is enough. It subtracts the target from the prediction, squares the error so misses in either direction are positive, and averages the result.

def mean_squared_error(
    prediction: torch.Tensor,
    target: torch.Tensor,
) -> torch.Tensor:
    error = prediction - target
    return torch.mean(error ** 2)

If the prediction is perfect, the loss is zero. If the prediction is wrong, the loss is positive.

target = torch.tensor([130.0, 210.0])
good_prediction = torch.tensor([130.0, 210.0])
bad_prediction = torch.tensor([100.0, 100.0])

print(mean_squared_error(good_prediction, target))
print(mean_squared_error(bad_prediction, target))

Output:

tensor(0.)
tensor(6500.)

The exact number is less important than the direction. Smaller is better, because this loss function says improvement means making predictions numerically closer to the targets.

A loss function is a contract about what improvement means.

If the loss rewards the wrong thing, training will faithfully optimize the wrong thing. That is not a PyTorch bug; that is a system design bug.

Gradients are not magic either

Now we can let PyTorch do the part it is very good at. Once the data, prediction, and loss are explicit, gradients have a clear job: tell us how changing each parameter would change the loss.

We create parameters and tell PyTorch to track their gradients:

hourly_rate = torch.tensor(1.0, requires_grad=True)
fixed_fee = torch.tensor(0.0, requires_grad=True)

Then we compute a prediction:

prediction = predict_cost(hours, hourly_rate, fixed_fee)

Then a loss:

loss = mean_squared_error(prediction, costs)

Then we ask PyTorch to work backward through the computation:

loss.backward()

After backward(), each parameter has a .grad. That gradient tells us how the loss changes when the parameter changes.

print(hourly_rate.grad)
print(fixed_fee.grad)

For training, we move each parameter slightly against that direction. The step is small because we want iterative improvement, not a blind jump.

with torch.no_grad():
    hourly_rate -= 0.01 * hourly_rate.grad
    fixed_fee -= 0.01 * fixed_fee.grad

The torch.no_grad() block matters because updating parameters is bookkeeping, not part of the model’s mathematical graph. We do not want PyTorch to track the update itself as another differentiable operation.

Then we clear the gradients:

hourly_rate.grad.zero_()
fixed_fee.grad.zero_()

This also matters because PyTorch accumulates gradients by default. That default is useful for some training patterns, but surprising at the beginning.

If you forget to clear gradients, the next update includes old gradient information. That gives us another small invariant:

Every training step should use the gradient for that step, not stale history.

The full training loop

Putting the pieces together gives us a small training loop. It predicts, measures loss, computes gradients, updates the parameters, clears the gradients, and repeats.

from __future__ import annotations

import torch


def predict_cost(
    hours: torch.Tensor,
    hourly_rate: torch.Tensor,
    fixed_fee: torch.Tensor,
) -> torch.Tensor:
    return hourly_rate * hours + fixed_fee


def mean_squared_error(
    prediction: torch.Tensor,
    target: torch.Tensor,
) -> torch.Tensor:
    error = prediction - target
    return torch.mean(error ** 2)


def train_linear_cost_model(
    hours: torch.Tensor,
    costs: torch.Tensor,
    *,
    steps: int = 1_500,
    lr: float = 0.01,
) -> tuple[torch.Tensor, torch.Tensor]:
    hourly_rate = torch.tensor(1.0, requires_grad=True)
    fixed_fee = torch.tensor(0.0, requires_grad=True)

    for _ in range(steps):
        prediction = predict_cost(hours, hourly_rate, fixed_fee)
        loss = mean_squared_error(prediction, costs)

        loss.backward()

        with torch.no_grad():
            hourly_rate -= lr * hourly_rate.grad
            fixed_fee -= lr * fixed_fee.grad

            hourly_rate.grad.zero_()
            fixed_fee.grad.zero_()

    return hourly_rate.detach(), fixed_fee.detach()

And use it:

hours = torch.tensor([2.0, 4.0, 6.0, 8.0])
costs = torch.tensor([130.0, 210.0, 290.0, 370.0])

hourly_rate, fixed_fee = train_linear_cost_model(hours, costs)

print(hourly_rate)
print(fixed_fee)
print(predict_cost(torch.tensor([10.0]), hourly_rate, fixed_fee))

The result should be close to:

tensor(40.)
tensor(50.)
tensor([450.])

The exact values may differ slightly depending on the learning rate and number of steps, and that is fine. Training is not symbolic algebra; it is iterative improvement under a contract.

The small loop is a changing surface: predict, measure, differentiate, update, repeat. 3D render by Logan Voss on Unsplash.

The same thing with torch.nn

The manual version is useful because it shows the pieces. But once those pieces are visible, PyTorch’s abstractions become less mysterious.

The same model can be written with torch.nn.Linear:

import torch
from torch import nn


model = nn.Linear(in_features=1, out_features=1)

This model expects two-dimensional input:

[number_of_examples, number_of_features]

So the hours tensor becomes:

hours = torch.tensor([[2.0], [4.0], [6.0], [8.0]])
costs = torch.tensor([[130.0], [210.0], [290.0], [370.0]])

The loss function can come from PyTorch too:

loss_fn = nn.MSELoss()

And the optimizer handles the parameter update:

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

The loop becomes shorter because nn.Linear owns the parameters and the optimizer knows how to update them.

for _ in range(1_500):
    prediction = model(hours)
    loss = loss_fn(prediction, costs)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Now inspect the learned parameters:

print(model.weight)
print(model.bias)

For this problem, the weight is the hourly rate and the bias is the fixed fee.

weight ~= 40
bias   ~= 50

The abstraction did not change the problem. It only removed bookkeeping, which is the right way to use abstractions: first understand what they replace, then let them replace it.

The shape problem appears immediately

The manual function accepted this:

hours = torch.tensor([2.0, 4.0, 6.0, 8.0])

But nn.Linear wants this:

hours = torch.tensor([[2.0], [4.0], [6.0], [8.0]])

The values are the same, but the shape is not.

print(torch.tensor([2.0, 4.0, 6.0, 8.0]).shape)
print(torch.tensor([[2.0], [4.0], [6.0], [8.0]]).shape)

Output:

torch.Size([4])
torch.Size([4, 1])

The second tensor says:

4 examples
1 feature per example

That is one of the most practical PyTorch lessons. Most beginner errors are not about calculus; they are about shape.

The model expected [batch, features], but the data was [batch]. Or the target was [batch] while the prediction was [batch, 1].

Sometimes PyTorch will complain, and sometimes it will broadcast and keep going. So I like writing shape checks near the boundary:

assert hours.ndim == 2
assert hours.shape[1] == 1
assert costs.shape == hours.shape

These checks are boring, and boring is good here. Boring protects the model from clever mistakes.

Most beginner PyTorch errors are shape errors wearing a math costume. 3D render by Logan Voss on Unsplash.

A slightly more realistic problem

Now add one more feature. Project cost is not only a function of hours; suppose urgent work has a premium.

The examples now look like:

hours urgent -> cost
2     0      -> 130
4     0      -> 210
6     1      -> 340
8     1      -> 420

The model is:

cost = hours_weight * hours + urgency_weight * urgent + fixed_fee

The input tensor now has two features per example. Each row is one project, and each column is one feature about that project.

x = torch.tensor(
    [
        [2.0, 0.0],
        [4.0, 0.0],
        [6.0, 1.0],
        [8.0, 1.0],
    ]
)

y = torch.tensor(
    [
        [130.0],
        [210.0],
        [340.0],
        [420.0],
    ]
)

The model changes by one number:

model = nn.Linear(in_features=2, out_features=1)

Everything else is the same:

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for _ in range(2_000):
    prediction = model(x)
    loss = loss_fn(prediction, y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

The learned weights now have a direct interpretation:

weight[0] -> cost per hour
weight[1] -> urgency premium
bias      -> fixed fee

The shape changed from:

[4, 1]

to:

[4, 2]

That is a bigger conceptual move than it looks. We are no longer feeding the model one number; we are feeding it a row of features.

This is where a lot of applied machine learning work lives:

What should be a feature?
What should not be a feature?
What does the target actually mean?
What would make this label dishonest?

PyTorch will optimize the parameters. It will not tell you whether the problem was framed honestly.

What PyTorch is doing for us

In this small exercise, PyTorch gives us four useful tools. None of them is mysterious once the problem is shaped clearly.

First, tensors:

Numbers with shape.

Second, automatic differentiation:

Compute gradients without manually deriving every formula.

Third, modules:

Objects that hold parameters and define forward computation.

Fourth, optimizers:

Bookkeeping for updating parameters from gradients.

That is enough to solve many basic problems, and it is also enough to understand larger systems more calmly. A neural network is not a spell; it is a function with parameters.

Training is not magic either. It is an update loop trying to reduce a loss, and the hard part is often not PyTorch but deciding what problem the loss is allowed to represent.

Complete listings to type through

The pieces above are easier to understand when they can be typed as one small program. These are deliberately complete rather than clever: each listing has its data, its training loop, its checks, and one prediction at the end.

Type the manual version first. It makes the gradient, update, and gradient reset visible. Then type the torch.nn version and notice that the numerical contract did not change; only the bookkeeping moved into PyTorch’s abstractions.

1. Manual autograd: one feature, two parameters

from __future__ import annotations

import torch


def predict_cost(
    hours: torch.Tensor,
    hourly_rate: torch.Tensor,
    fixed_fee: torch.Tensor,
) -> torch.Tensor:
    return hourly_rate * hours + fixed_fee


def mean_squared_error(
    prediction: torch.Tensor,
    target: torch.Tensor,
) -> torch.Tensor:
    return torch.mean((prediction - target) ** 2)


def train_linear_cost_model(
    hours: torch.Tensor,
    costs: torch.Tensor,
    *,
    steps: int = 1_500,
    lr: float = 0.01,
) -> tuple[torch.Tensor, torch.Tensor]:
    hourly_rate = torch.tensor(1.0, dtype=hours.dtype, requires_grad=True)
    fixed_fee = torch.tensor(0.0, dtype=hours.dtype, requires_grad=True)

    for _ in range(steps):
        prediction = predict_cost(hours, hourly_rate, fixed_fee)
        loss = mean_squared_error(prediction, costs)
        loss.backward()

        with torch.no_grad():
            hourly_rate -= lr * hourly_rate.grad
            fixed_fee -= lr * fixed_fee.grad
            hourly_rate.grad.zero_()
            fixed_fee.grad.zero_()

    return hourly_rate.detach(), fixed_fee.detach()


hours = torch.tensor([2.0, 4.0, 6.0, 8.0])
costs = torch.tensor([130.0, 210.0, 290.0, 370.0])

hourly_rate, fixed_fee = train_linear_cost_model(hours, costs)
ten_hour_cost = predict_cost(torch.tensor([10.0]), hourly_rate, fixed_fee)

assert torch.isclose(hourly_rate, torch.tensor(40.0), atol=0.5)
assert torch.isclose(fixed_fee, torch.tensor(50.0), atol=1.0)
assert torch.isclose(ten_hour_cost, torch.tensor([450.0]), atol=2.0).all()

print(f"hourly rate: {hourly_rate.item():.2f}")
print(f"fixed fee: {fixed_fee.item():.2f}")
print(f"10-hour cost: {ten_hour_cost.item():.2f}")

The important line to pause over is not loss.backward(). It is the pair of parameter updates inside torch.no_grad(), followed by the reset. Those three lines are the manual form of what the optimizer will take over later.

2. torch.nn: the same one-feature contract

import torch
from torch import nn


torch.manual_seed(0)

hours = torch.tensor([[2.0], [4.0], [6.0], [8.0]])
costs = torch.tensor([[130.0], [210.0], [290.0], [370.0]])

model = nn.Linear(in_features=1, out_features=1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for _ in range(1_500):
    prediction = model(hours)
    loss = loss_fn(prediction, costs)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

hourly_rate = model.weight.detach().item()
fixed_fee = model.bias.detach().item()
ten_hour_cost = model(torch.tensor([[10.0]])).item()

assert abs(hourly_rate - 40.0) < 0.5
assert abs(fixed_fee - 50.0) < 1.0
assert abs(ten_hour_cost - 450.0) < 2.0

print(f"hourly rate: {hourly_rate:.2f}")
print(f"fixed fee: {fixed_fee:.2f}")
print(f"10-hour cost: {ten_hour_cost:.2f}")

Here, model.parameters() gives the optimizer the weight and bias that the manual example named explicitly. optimizer.zero_grad(), loss.backward(), and optimizer.step() are the same rhythm, compressed into the conventional PyTorch interface.

3. Two features: hours and urgency

import torch
from torch import nn


torch.manual_seed(0)

# Each row is [hours, urgent].
x = torch.tensor(
    [
        [2.0, 0.0],
        [4.0, 0.0],
        [6.0, 1.0],
        [8.0, 1.0],
    ]
)
y = torch.tensor([[130.0], [210.0], [340.0], [420.0]])

model = nn.Linear(in_features=2, out_features=1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for _ in range(2_000):
    prediction = model(x)
    loss = loss_fn(prediction, y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

hours_weight, urgency_weight = model.weight.detach().flatten().tolist()
fixed_fee = model.bias.detach().item()
urgent_ten_hour_cost = model(torch.tensor([[10.0, 1.0]])).item()

assert abs(hours_weight - 40.0) < 0.5
assert abs(urgency_weight - 50.0) < 1.0
assert abs(fixed_fee - 50.0) < 1.0
assert abs(urgent_ten_hour_cost - 500.0) < 2.0

print(f"hours weight: {hours_weight:.2f}")
print(f"urgency weight: {urgency_weight:.2f}")
print(f"fixed fee: {fixed_fee:.2f}")
print(f"urgent 10-hour cost: {urgent_ten_hour_cost:.2f}")

This last listing is the useful transition point. The training loop barely changes, while the meaning of one row becomes richer. That is the pattern to carry forward: make the input shape say what one example means, then make the loss say what a good answer means.

The small checklist I would keep

When learning PyTorch, I would keep the checklist short. The questions are simple, but they catch most of the places where a small model becomes confusing.

What is the input?
What is the target?
What is the shape of each tensor?
What are the parameters?
What does prediction mean?
What does loss mean?
Can the loss go down for the wrong reason?
Can I test one tiny example by hand?

If those questions are clear, the code is usually manageable. If those questions are unclear, adding a bigger model usually makes the confusion more expensive.

So the first PyTorch skill is not writing a neural network. It is turning a problem into a small numerical contract, so tensors, gradients, and optimizers have something honest to do.

Part of the "primitives first, then compositions" series.

Need a small PyTorch prototype before the model gets large?

I agree to the Privacy Policy Shape the Experiment