How to implement a custom operator that support multiple compute device (CPU, CUDA)?

[wangxianliang]

Ask a Question

I tried the following implementation, but had no effect.

CUDA implementation:

struct CustomOPGpu : Ort::CustomOpBase<CustomOPGpu , CustomKernel> {
		const char* GetName() const { return "CustomOP"; };
const char* GetExecutionProviderType() const { return "CUDAExecutionProvider"; };
...
}

CPU implementation:

struct CustomOPCpu : Ort::CustomOpBase<CustomOPCpu , CustomKernel> {
		const char* GetName() const { return "CustomOP"; };
const char* GetExecutionProviderType() const { return "CPUExecutionProvider"; };
...
}

The doc (https://onnxruntime.ai/docs/reference/operators/add-custom-op.html) doesn't have any sample codes.

Question

Further information

• Relevant Area:

• Is this issue related to a specific model?
  Model name:
  Model opset:

Notes