How to pass Python device and dtype to C++

I’m trying to implement a C++ function that has the arguments of a Device and Dtype and will be called in Python. But I got error trying to achieve this.
Here’s my C++ code:

#define tensor torch::Tensorv
tensor fun1(tensor x,torch::Device device,torch::Dtype dtype) {
	auto opt = torch::TensorOptions(device).dtype(dtype);

Here’s the Python part:


Here’s the error:

TypeError: fun1(): incompatible function arguments. The following argument types are supported:
    1. (arg0: at::Tensor, arg1: c10::Device, arg2: c10::ScalarType) -> at::Tensor

Invoked with: tensor([1]), device(type='cpu'), torch.float32

I don’t think this works directly right now.

If you have compiled torch yourself, you see that in torch/csrc/autograd/generated/python_torch_functions.cpp the factory functions take the Python objects and unpacks them.

The PyBind converters are only present for IntArrayRef and Tensor (toch/csrc/utils/pybind.h). You could try to register your own for torch::Device and torch::Dtype – or submit that PR (1.5 years ago, it was only Tensors).

A standard workaround could be to just pass a (potentially empty) tensor and take the device and dtype of that.

Best regards


Ok, thank you very much. I’ll try that.

Actually, I found a way to pass dtype to C++.
In TouchScript(JIT), torch.dtype type are treated as python int, which corressponds to what happens in C++, as the dtypes are all ScalarType. So I can just create an C++ function that takes in an integer and make it an ScalarType to pass the dtype. But this method can only be called in TouchScript because only in TouchScript are dtypes treated as python int.