How to check the size of tensor in cuda extension

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

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

The full API can be found in the headers…

Thank you very much!