PyTorch First Program

Writing your first PyTorch program is the fastest way to understand how PyTorch works in practice. This page walks you through three short programs — each one builds on the last — so you leave with a clear picture of PyTorch's core flow.

Program 1: Hello Tensor

A tensor is the fundamental unit of data in PyTorch — think of it as a container that holds numbers. In your first program, you create a tensor and print it.

import torch

x = torch.tensor([1.0, 2.0, 3.0])
print(x)

Output:

tensor([1., 2., 3.])

You just stored three numbers in a PyTorch tensor. The word tensor in the output tells you the type. The numbers inside are the values you stored. That's it — your first PyTorch object.

Program 2: Simple Math with Tensors

Tensors support all standard math operations. You can add, subtract, multiply, and divide them the same way you would with plain Python numbers.

import torch

a = torch.tensor([10.0, 20.0, 30.0])
b = torch.tensor([1.0, 2.0, 3.0])

print(a + b)
print(a * b)

Output:

tensor([11., 22., 33.])
tensor([ 10.,  40.,  90.])

Notice that the operations happen element-by-element. The first element of a (which is 10) gets added to the first element of b (which is 1), giving 11. This parallel element-wise behavior is what makes tensors so efficient for neural network math.

Diagram: How Element-Wise Addition Works

Tensor a:  [ 10,  20,  30 ]
               |    |    |
Tensor b:  [  1,   2,   3 ]
               ↓    ↓    ↓
Result:    [ 11,  22,  33 ]

Program 3: Your First Mini Neural Network

This program builds the simplest possible neural network — one that takes a number as input and produces a number as output — then trains it to learn a pattern.

The pattern to learn: output = input × 2

import torch
import torch.nn as nn

# Step 1: Define a simple model
model = nn.Linear(1, 1)  # 1 input, 1 output

# Step 2: Define a loss function and optimizer
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Step 3: Create training data
X = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[2.0], [4.0], [6.0], [8.0]])

# Step 4: Train for 100 rounds
for epoch in range(100):
    pred = model(X)
    loss = loss_fn(pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Step 5: Test the model
test_input = torch.tensor([[5.0]])
print(model(test_input))

Output (approximate):

tensor([[9.9998]], grad_fn=<AddmmBackward0>)

The model predicted roughly 10 when given 5 — very close to the correct answer of 10 (since 5 × 2 = 10). The model learned the pattern by adjusting its internal values 100 times based on the training data you provided.

Understanding Each Step

Step 1 – Define the Model

nn.Linear(1, 1) creates a single layer with one input and one output. Internally, it holds a weight and a bias — two numbers it adjusts during training to fit the data.

Step 2 – Loss Function and Optimizer

The loss function (MSELoss) measures how far the model's prediction is from the correct answer. The optimizer (SGD) adjusts the model's weight and bias to reduce that gap.

Step 3 – Training Data

You gave the model four examples. For input 1, the correct output is 2. For input 2, the correct output is 4. The model uses these pairs to learn the rule.

Step 4 – The Training Loop

Each epoch (round):
  ┌──────────────────────────────────┐
  │ 1. model makes prediction        │
  │ 2. loss measures the error       │
  │ 3. zero_grad clears old values   │
  │ 4. backward computes gradients   │
  │ 5. step adjusts the weights      │
  └──────────────────────────────────┘
              ↓ repeat 100 times

Step 5 – Test the Model

After training, you fed the model a new number it had never seen before (5.0). It correctly predicted approximately 10, which shows it learned the rule and didn't just memorize the training data.

Reading the Output

The output tensor([[9.9998]], grad_fn=<AddmmBackward0>) might look strange. The double brackets appear because the tensor has a 2D shape. The grad_fn part shows that PyTorch is tracking this value for gradient computation — you can ignore it for now.

What Just Happened — The Big Picture

Training Data  →  Model  →  Prediction
                    ↑
               Optimizer adjusts
                    ↑
               Loss Function measures error
                    ↑
               Correct Answer

This loop — predict, measure error, adjust — is the engine behind every neural network, from a simple two-number model to a billion-parameter language model. You just ran that engine for the first time.

Summary

Your first PyTorch program created a tensor and printed it. Your second performed element-wise math. Your third built and trained a tiny neural network that learned to multiply numbers by 2. Every PyTorch project, no matter how complex, follows the same core steps: define a model, define a loss function and optimizer, feed it data, and run the training loop. You now understand that foundation.