error with loading torchscript model

[viktortnk]

Hello, I'm getting error with loading compiled model with torch.jit.load
Interesting that it works within python scripts codebase used to train the model, but doesn't work from python interpreter and C++ code. it is reporting the same error

Ultimately I need the model loaded in C++ for inference.

Could anyone suggest how to possibly solve this?
thanks

>>> new_model = torch.jit.load(export_path)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3/dist-packages/torch/jit/_serialization.py", line 162, in load
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files)
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 6
  _0 = __torch__.torch_geometric._compile.is_compiling
  _1 = __torch__.torch_geometric.utils._segment._torch_segment
  _2 = __torch__.torch_scatter.segment_csr.segment_csr
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  if _0():
    _3 = _1(src, ptr, reduce, )
'segment' is being compiled since it was called from 'SumAggregation.reduce'
Serialized   File "code/__torch__/torch_geometric/nn/aggr/basic.py", line 23
    _1 = __torch__.torch_geometric._compile.is_compiling
    _2 = __torch__.torch_geometric.nn.aggr.base.expand_left
    _3 = __torch__.torch_geometric.utils._segment.segment
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _4 = "Aggregation requires \'index\' to be specified"
    _5 = __torch__.torch_geometric.utils._scatter.scatter
'SumAggregation.reduce' is being compiled since it was called from 'SumAggregation.forward'
  File "/home/vltaranenko/.local/lib/python3.10/site-packages/torch_geometric/nn/aggr/basic.py", line 22
                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,

[rusty1s]

Sorry for late reply. Do you have a minimal script to reproduce?