Adding shape inference to custom operator for ONNX exporting

The warning disappeared with .setType and torch==1.14.0.dev20221121+cu116