🎓 A1Learner
🔍
0 of 79 lessons completed

PyTorch Indexing Tensors

Indexing picks one part of a tensor by its position. One index picks a row, and two indexes pick one cell in a table.

Pick a whole row

One index picks one row from a table.

scores = torch.tensor([[85, 92, 78], [90, 88, 95], [72, 81, 84]])
print(scores[0])

Pick one cell

Two indexes name one cell. The first number is the row, and the second is the column.

print(scores[0, 2])
print(scores[2, 1])

Count from the end

A negative index counts back from the end.

print(scores[-1])
print(scores[-1, -1])

Example

import torch

scores = torch.tensor([[85, 92, 78],
                       [90, 88, 95],
                       [72, 81, 84]])

print("Table shape:", scores.shape)
print("Ana's scores:", scores[0])
print("Ana's third score:", scores[0, 2])
print("Ben's last score:", scores[1, -1])
print("Cara's second score:", scores[2, 1])
print("Last row:", scores[-1])

Key points

  • Index: a number that gives the position of a row or a cell
  • Use scores[0] to get the first row of a table
  • Use scores[1, 2] to get the cell in row 1 and column 2
  • The first position is 0, not 1
  • A negative index counts from the end, so -1 is the last position

In this example, scores is a table with three rows and three columns. The first print shows the table shape. scores[0] picks Ana's row, the first row. scores[0, 2] picks the cell in row 0 and column 2. The next print uses -1 to pick Ben's last score. scores[2, 1] picks Cara's second score. The last print uses -1 alone to pick the last row.

Note

A single cell still prints as a tensor, like tensor(78)

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does scores[1, 0] give from a table of scores?