How to get a column from a tensor?

user_13855593 0
Viewed 1782

Let's say I have a tensor consisting of 1 and 0's as shown below. How can I get the index of a specific column to replace with new values ? If I want to replace the values of column 1 with the [3.,4.,5.,6.], how do I accomplish this ?

a = torch.tensor([[[1., 0., 0., 0.]],
        [[0., 1., 0., 0.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]]])
tensor pytorch
1 Answers

Calling them 'columns' is a bit tricky, given that this is a 3D tensor.

This will do what you need, setting 'column' 1 to the values you gave.

a = torch.tensor([[[1., 0., 0., 0.]],
        [[0., 1., 0., 0.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]]])

# Change values in 'column' 1 (zero-indexed):
# The 0 is there because of the size-1 second dimension.
a[1, 0, :] = torch.tensor([3., 4., 5., 6.])

print(a)
# tensor([[[1., 0., 0., 0.]],
#         [[3., 4., 5., 6.]],
#         [[1., 0., 0., 0.]],
#         [[0., 0., 0., 1.]],
#         [[1., 0., 0., 0.]],
#         [[0., 0., 0., 1.]],
#         [[1., 0., 0., 0.]]])
user_78022000
Related