How to check the size of tensor in cuda extension

1

I am building a cuda extension following the PyTorch C FFI examples. When we pass the data to C function in THCudaTensor*, is there any method that we can check the size of each dimension of the data or we have to pass the batch_size, nchannels, width and height simultaneously? Thank you very much for the help in advance.

1 Like
2

you can use: THCTensor_(size)(input, dimension)

The full API can be found in the headers…
https://github.com/pytorch/pytorch/blob/master/torch/lib/THC/generic/THCTensor.h#L26

3

Thank you very much!