Search⌘ K
Home/
Courses/
Basic Pytorch Tensor Manipulation for Machine Learning

Selecting Elements from a Tensor

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

We'll cover the following...

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, :].

Python 3.8
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: reshape returns a ...