ONNX export failed on an operator with unrecognized namespace 'torch_scatter::scatter_max'

I have a pytorch network like this

import torch.nn as nn
import torch_scatter.scatter_max

class DGCN(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        ...
        torch_scatter.scatter_max(x, index, dim=0)
        ...

But when i want export my model to onnx, i face this error:

  ...
  File "/usr/local/lib/python3.9/dist-packages/torch/onnx/utils.py", line 1115, in _model_to_graph
    graph = _optimize_graph(
  File "/usr/local/lib/python3.9/dist-packages/torch/onnx/utils.py", line 663, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "/usr/local/lib/python3.9/dist-packages/torch/onnx/utils.py", line 1909, in _run_symbolic_function
    raise errors.UnsupportedOperatorError(
torch.onnx.errors.UnsupportedOperatorError: ONNX export failed on an operator with unrecognized namespace 'torch_scatter::scatter_max'. 
If you are trying to export a custom operator, make sure you registered it with the right domain and version.

So, How i can do this exactly?

This issue was due to the torch_scatter::scatter_max not being supported by either stock PyTorch or ONNX.

Some users have the same issue.

To solve this issue, I think you can implement the torch_scatter::scatter_max by torch.nn.Module or TORCH.NN.FUNCTIONAL.