Indexing and slicing in PyTorch

Introduction

Indexing and slicing help you pick parts of data from tensors easily, like choosing pieces from a big cake.

When you want to get a single value from a tensor.

When you need a smaller part of a tensor for calculations.

When you want to change some values inside a tensor.

When you want to loop over parts of data in a tensor.

When you want to prepare data batches for training a model.

Syntax

PyTorch

tensor[start:stop:step]
tensor[index]
tensor[start:stop]
tensor[:, column_index]
tensor[row_index, :]

Use ':' to select ranges, like start to stop (stop not included).

You can use commas to index multiple dimensions, like tensor[row, column].

Examples

Gets the second element (index 1) from the 1D tensor.

PyTorch

x = torch.tensor([10, 20, 30, 40, 50])
print(x[1])

Gets elements from index 1 up to but not including 4.

PyTorch

x = torch.tensor([10, 20, 30, 40, 50])
print(x[1:4])

Gets all rows but only the second column.

PyTorch

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(x[:, 1])

Gets the first row and first two columns.

PyTorch

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(x[0, :2])

Sample Model

This program shows how to pick single elements, rows, columns, and smaller parts from a 2D tensor using indexing and slicing.

PyTorch

import torch

# Create a 2D tensor
x = torch.tensor([[10, 20, 30], [40, 50, 60], [70, 80, 90]])

# Get the element at row 1, column 2
elem = x[1, 2]

# Get the first row
first_row = x[0, :]

# Get the last column
last_col = x[:, 2]

# Get a sub-tensor of first two rows and first two columns
sub_tensor = x[:2, :2]

print(f"Element at (1, 2): {elem}")
print(f"First row: {first_row}")
print(f"Last column: {last_col}")
print(f"Sub-tensor (2x2):\n{sub_tensor}")

Important Notes

Remember that indexing starts at 0, so the first element is at index 0.

Negative indices count from the end, so -1 is the last element.

Slicing does not include the stop index, so x[0:2] gets elements 0 and 1.

Summary

Indexing picks single elements by position.

Slicing picks ranges of elements using start:stop:step.

Use commas to index multiple dimensions in tensors.