PyTorch Tensor Indexing

Indexing lets you access specific elements, rows, columns, or slices inside a tensor. PyTorch indexing follows the same rules as Python list indexing and NumPy — which means if you know one, you already know the basics of the others.

Basic Indexing

Use square brackets with a position number to access elements. Positions start at 0. Negative numbers count from the end.

import torch

t = torch.tensor([10, 20, 30, 40, 50])

print(t[0])    # tensor(10)   ← first element
print(t[-1])   # tensor(50)   ← last element
print(t[2])    # tensor(30)   ← third element

Indexing a 2D Tensor

For a 2D tensor, provide two indices: [row, column]. You can also use separate brackets: [row][column].

t = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])

print(t[0, 0])   # tensor(1)   ← row 0, col 0
print(t[1, 2])   # tensor(6)   ← row 1, col 2
print(t[2, -1])  # tensor(9)   ← row 2, last column

t =  [[ 1,  2,  3],     row 0
      [ 4,  5,  6],     row 1
      [ 7,  8,  9]]     row 2
        ↑   ↑   ↑
       col0 col1 col2

t[1, 2] → row 1, col 2 → 6

Slicing

Slicing selects a range of elements using the format start:stop:step. The stop index is excluded.

t = torch.tensor([0, 10, 20, 30, 40, 50])

print(t[1:4])    # tensor([10, 20, 30])   ← positions 1, 2, 3
print(t[::2])    # tensor([ 0, 20, 40])   ← every 2nd element
print(t[:3])     # tensor([ 0, 10, 20])   ← first 3
print(t[3:])     # tensor([30, 40, 50])   ← from position 3 onward

Slicing 2D Tensors

Apply slicing separately to rows and columns using a comma.

t = torch.tensor([[1, 2, 3, 4],
                  [5, 6, 7, 8],
                  [9, 10, 11, 12]])

print(t[:, 0])      # all rows, column 0 → tensor([1, 5, 9])
print(t[0, :])      # row 0, all columns → tensor([1, 2, 3, 4])
print(t[1:3, 1:3])  # rows 1-2, cols 1-2 → 2x2 sub-matrix

t[:, 0] — select entire first column:
[[ 1 ← ,  2,  3,  4],
 [ 5 ← ,  6,  7,  8],
 [ 9 ← , 10, 11, 12]]
Result: [1, 5, 9]

Boolean Indexing (Masking)

Pass a boolean tensor as an index to select only elements where the mask is True. This is extremely useful for filtering data.

t = torch.tensor([3.0, -1.0, 7.0, -5.0, 2.0])

mask = t > 0
print(mask)       # tensor([ True, False,  True, False,  True])
print(t[mask])    # tensor([3., 7., 2.])  ← only positive values

t:     [ 3,  -1,   7,  -5,   2]
mask:  [ T,   F,   T,   F,   T]
         ↓         ↓         ↓
select: [ 3,       7,        2]

Fancy Indexing

You can pass a list of specific indices to select multiple elements at arbitrary positions — not just a contiguous slice.

t = torch.tensor([100, 200, 300, 400, 500])
idx = torch.tensor([0, 2, 4])    # pick positions 0, 2, 4
print(t[idx])   # tensor([100, 300, 500])

Indexing 3D Tensors

A 3D tensor is like a stack of 2D tables. Index with three values: [depth, row, column].

t = torch.rand(4, 3, 5)   # 4 layers, 3 rows, 5 columns

print(t[0].shape)          # torch.Size([3, 5])  ← first layer
print(t[0, 1].shape)       # torch.Size([5])     ← first layer, second row
print(t[0, 1, 2])          # single element

3D tensor shape [4, 3, 5] — think of it as:
4 pages, each containing a 3×5 table

t[0]       → entire page 0 (shape 3×5)
t[0, 1]    → page 0, row 1 (shape 5)
t[0, 1, 2] → page 0, row 1, column 2 (scalar)

Modifying Elements via Indexing

You can assign new values to specific positions in a tensor using the same indexing syntax.

t = torch.zeros(3, 3)
t[1, 1] = 99.0     # set center element to 99
t[:, 2] = 7.0      # set entire last column to 7
print(t)

torch.where – Conditional Selection

torch.where picks values from two tensors based on a condition — element-by-element. Think of it as a ternary operator applied to entire tensors.

a = torch.tensor([1.0, 2.0, 3.0, 4.0])
b = torch.tensor([10.0, 10.0, 10.0, 10.0])
condition = a > 2

result = torch.where(condition, a, b)
print(result)   # tensor([10., 10.,  3.,  4.])
# where a > 2, keep a; otherwise use b

Summary

PyTorch indexing uses Python-style bracket notation. Single integers select specific elements. Negative indices count from the end. Slices (start:stop:step) select ranges. For 2D tensors, combine row and column indices with a comma. Boolean masks filter elements based on conditions. Fancy indexing picks elements at arbitrary positions. torch.where conditionally selects between two tensors. These techniques appear constantly when preprocessing data, inspecting model outputs, and building custom loss functions.