How to select columns from 2D tensor

1

I have a DoubleTensor A of size 4x4 and I want to make a new tensor B with the first k columns of A.

A:
-0.5961 -0.2435 0.0110 -1.1665
-0.4133 -0.6013 0.1223 1.5625
-0.1788 0.4041 1.2526 0.0308
-0.0207 0.9357 -1.2322 -0.5882
[torch.DoubleTensor of size 4x4]

Desired B if k = 2:
-0.5961 -0.2435
-0.4133 -0.6013
-0.1788 0.4041
-0.0207 0.9357
[torch.DoubleTensor of size 4x2]

This is my current solution:
k = 2
cols = torch.chunk(A, 4, 1)
cols = cols[:k]
B = torch.cat(cols, 1)

Is there a better way to do this?

3 Likes
2

PyTorch supports numpy-style indexing. To select the first 2 columns

a = torch.rand(4,4)
a[:, :2]

should do the trick.

7 Likes
4

you can do
B = A[ row_start : row_end , column_start:column_end].clone()

for your example use
B = A[: , 0:k].clone()

2 Likes