Problem when loading a jit model in your environment

[kizym]

Hi, I have followed your tutorial on how to install the project locally from source. now I have a model already trained in another environment and when I try to load it in your environment I get the following error

Traceback (most recent call last):
  File "/Users/kizym/torch-mlir/examples/model/torch_model.py", line 3, in <module>
    load = torch.jit.load("model.pt")
  File "/Users/kizym/mlir_venv/lib/python3.10/site-packages/torch/jit/_serialization.py", line 162, in load
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files, _restore_shapes)  # type: ignore[call-arg]
RuntimeError: 
Unknown builtin op: torch_scatter::segment_sum_csr.
Could not find any similar ops to torch_scatter::segment_sum_csr. This op may not exist or may not be currently supported in TorchScript.
:
  File "code/__torch__/torch_scatter/segment_csr.py", line 35
    indptr: Tensor,
    out: Optional[Tensor]=None) -> Tensor:
  _10 = ops.torch_scatter.segment_sum_csr(src, indptr, out)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _10
def segment_mean_csr(src: Tensor,
'segment_sum_csr' is being compiled since it was called from 'segment_csr'
Serialized   File "code/__torch__/torch_scatter/segment_csr.py", line 5
    out: Optional[Tensor]=None,
    reduce: str="sum") -> Tensor:
  _0 = __torch__.torch_scatter.segment_csr.segment_sum_csr
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  _1 = __torch__.torch_scatter.segment_csr.segment_mean_csr
  _2 = __torch__.torch_scatter.segment_csr.segment_min_csr
'segment_csr' is being compiled since it was called from 'segment'
Serialized   File "code/__torch__/torch_geometric/utils/segment.py", line 4
    ptr: Tensor,
    reduce: str="sum") -> Tensor:
  _0 = __torch__.torch_scatter.segment_csr.segment_csr
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _0(src, ptr, None, reduce, )
'segment' is being compiled since it was called from 'SumAggregation.reduce'
Serialized   File "code/__torch__/torch_geometric/nn/aggr/basic.py", line 22
    reduce: str="sum") -> Tensor:
    _1 = __torch__.torch_geometric.nn.aggr.base.expand_left
    _2 = __torch__.torch_geometric.utils.segment.segment
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _3 = __torch__.torch_geometric.utils.scatter.scatter
    _4 = uninitialized(Tensor)
'SumAggregation.reduce' is being compiled since it was called from 'SumAggregation.forward'
  File "/Users/kizym/miniconda3/envs/pytorch/lib/python3.10/site-packages/torch/nn/aggr/basic.py", line 21
                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                dim: int = -2) -> Tensor:
        return self.reduce(x, index, ptr, dim_size, dim, reduce='sum')
                                                                ~~~~ <--- HERE
Serialized   File "code/__torch__/torch_geometric/nn/aggr/basic.py", line 12
    dim_size: Optional[int]=None,
    dim: int=-2) -> Tensor:
    _0 = (self).reduce(x, index, ptr, dim_size, dim, "sum", )
                                                     ~~~~~ <--- HERE
    return _0
  def reduce(self: __torch__.torch_geometric.nn.aggr.basic.SumAggregation,

This is the simple file I am using

import torch

load = torch.jit.load("model.pt")

by looking on the web I found this issue rusty1s/pytorch_scatter#147 (comment) which says that I should build the c++ API. but how can I do it in your environment ? The provided tutorials are not considering this case, and the model is loaded without problems in the environment from which it has been generated