Selecting Elements from a Tensor

In this lesson, we look at selecting elements from a tensor.

Selecting tensor with index

If you are already familiar with the NumPy array, you can use the same methods to select tensors by [] operations.

Take a 2-dimensional tensor as an example. Let’s consider it as a matrix.

  • tensor[2, 3]: Get only one value.
  • tensor[:, 1]: Get the second column from the tensor.
  • tensor[1, :]: Get the second row from the tensor.

For higher-dimensional tensors, the operations are the same. Such as tensor[:, 2, :].

Press + to interact
import torch
a = torch.arange(1, 10).reshape((3, 3))
# The output
# tensor([[1, 2, 3],
# [4, 5, 6],
# [7, 8, 9]])
print("The original tensor")
print(a)
print("Select only one element")
print(a[1,1])
print("Select the second column")
print(a[:, 1])
print("Select the second row.")
print(a[1, :])

Notice: ...

Get hands-on with 1400+ tech skills courses.