Runtime error, op 'tensordot' not implemented.

[yuxiaohui78]

🐞Describing the bug

My model uses the function torch.tensordot() and I tried to convert my model to CoreML model.
But I got the error below:

RuntimeError: PyTorch convert function for op 'tensordot' not implemented.

Could you please support this function in the next version of the coremltools?

Stack Trace

Converting PyTorch Frontend ==> MIL Ops:  86%|████████████████████████████████████████████████████████████████████████████████▌             | 6/7 [00:00<00:00, 6419.85 ops/s]
Traceback (most recent call last):
  File "/data/Model_convert/tensordot_test/test_tensordot.py", line 54, in <module>
    main()
  File "/data/Model_convert/tensordot_test/test_tensordot.py", line 44, in main
    mlmodel = ct.convert(
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/_converters_entry.py", line 574, in convert
    mlmodel = mil_convert(
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 188, in mil_convert
    return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 212, in _mil_convert
    proto, mil_program = mil_convert_to_proto(
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 286, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 108, in __call__
    return load(*args, **kwargs)
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 80, in load
    return _perform_torch_convert(converter, debug)
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 107, in _perform_torch_convert
    raise e
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 99, in _perform_torch_convert
    prog = converter.convert()
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 519, in convert
    convert_nodes(self.context, self.graph)
  File "/home/yxh/anaconda3/envs/DSV2/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 86, in convert_nodes
    raise RuntimeError(
RuntimeError: PyTorch convert function for op 'tensordot' not implemented.

To Reproduce

• Please add a minimal code example that can reproduce the error when running it.

import torch 
import torch.nn as nn
import numpy as np
import coremltools as ct  


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

    def forward(self, a, b):
        c = torch.tensordot(a, b, dims=([1, 0], [0, 1]))
        return c
               
def main():
    a = torch.arange(60.).reshape(3, 4, 5)
    
    b = torch.arange(24.).reshape(4, 3, 2)
    
    print (a)
    print (b)
    c = torch.tensordot(a, b, dims=([1, 0], [0, 1]))
    print (c)

    print ("=======test model=====")
    model = Net().to(torch.device("cpu"))
   
    model.eval()
    
    output = model(a, b)
    print (output)

    print ("========convert to coreml model=======")
    trace = torch.jit.trace(model, (a, b))
    #trace = torch.jit.script(model, (a, b))
    input_a = ct.TensorType(
           name='a',
           shape=a.shape
       )
    input_b = ct.TensorType(
           name='b',
           shape=b.shape
       )   
    mlmodel = ct.convert(
           trace,
           inputs=[input_a, input_b],
       )

    mlmodel.save("model.mlpackage")
    
    print ("Coreml model saved...")
        
if __name__ == '__main__':
    main()

• If the model conversion succeeds, but there is a numerical mismatch in predictions, please include the code used for comparisons.

System environment (please complete the following information):

coremltools version: 7.1
OS (e.g. MacOS version or Linux type): Ubuntu 20.04.4 LTS
Any other relevant version information (e.g. PyTorch or TensorFlow version): Pytorch 1.13.1+cu116

Additional context

• Add anything else about the problem here that you want to share.