Tensor manipulation in ATen C++

My problem is : I have a tensor with shape output1.shape torch.Size([1, 2, 120, 89]) .

I wanted to do something like output1.data.numpy()[0,1,:,:]. I did

auto output3=output1.squeeze() which produced output3.shape torch.Size([ 2, 120, 89]).

Now, how can I get [0,:,:] and [1,:,:] separately as 2 tensors each with shape [120,89] ?

I can use auto accessor=output3.accessor<float,3>(); and get the elements only by one as a 2D array.

Is there a better way to do these type of operations in C++ ?

Best Regards


May be you can try torch::slice().

Tensor select(int64_t dim, int64_t index) const (https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/Tensor.h#L458) also works.

Thanks for the suggestion.
I have another issue. I can run the model inference using C++ , underlying library is ATen. Right now it is running in One GPU. How can I make it run on multiple GPU ? Something similar to dataparallel in PyTorch.

Best regards