'Tensor.data_ptr()' not visible from TorchScript

jit
1

Hi all,

When I try to compile a TorchScript module that includes use of the spconv library, I get an error claiming that:

‘Tensor’ boject has no attribute or method ‘data_ptr’.

I believe this is because ‘data_ptr’ is defined in a parent class of Tensor - ‘TensorBase’. Since TorchScript doesn’t do inheritence, it can’t find this method.

I’ve tried asserting that the tensor is an instance of torch.BaseTensor (or torch._C.BaseTensor). That fails, claiming that that type can’t be found.

Can anyone think of a way to make this work?

2

I found a solution.

You can use torch.library.custom_op to register the operation. It does not show up by default because it is a call into C++.

Example:

@torch.library.custom_op("mylib::data_ptr", mutates_args=())
def data_ptr(tensor: torch.Tensor) -> int:
    return tensor.data_ptr()

@data_ptr.register_fake
def _(_: torch.Tensor) -> int:
    return 0

You then change every instance of tensor.data_ptr() to data_ptr(tensor).