I’m a nubie to python and pytorch with a very basic question.

I have a tensor T and I want to access column i of this tensor:

T[:,i]

when T is multidimensional (e.g., torch.Size([3, 4])) this works fine, but if it is not (e.g., torch.Size([4])), I get the following error:

IndexError: too many indices for tensor of dimension 1

In the latter case I wish to just access element i of the row vector. My question is: what code should I write so that my indexing can deal with both the M x N shaped tensor and the 1 X N shaped tensor.

Thanks in advance.