Tensor slicing issue: 'NoneType' object has no attribute 'shape'

[ryan6310]

Hi,

I am having an issue converting a simple torch model to coreml.

Trace

Converting Frontend ==> MIL Ops:  86%|██████████████████████████████████████████████████████████████▌          | 6/7 [00:00<00:00, 1074.31 ops/s]
Traceback (most recent call last):
  File "cm_test.py", line 29, in <module>
    model = ct.convert(
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/_converters_entry.py", line 352, in convert
    mlmodel = mil_convert(
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 183, in mil_convert
    return _mil_convert(model, convert_from, convert_to, ConverterRegistry, MLModel, compute_units, **kwargs)
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 210, in _mil_convert
    proto, mil_program = mil_convert_to_proto(
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 273, in mil_convert_to_proto
    prog = frontend_converter(model, **kwargs)
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/converter.py", line 105, in __call__
    return load(*args, **kwargs)
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 47, in load
    return _perform_torch_convert(converter, debug)
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 84, in _perform_torch_convert
    prog = converter.convert()
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 250, in convert
    convert_nodes(self.context, self.graph)
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 89, in convert_nodes
    add_op(context, node)
  File "/Users/youngre/miniconda3/envs/coreml3/lib/python3.8/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 2674, in select
    assert index.shape == ()
AttributeError: 'NoneType' object has no attribute 'shape'

To reproduce

  import torch
  import coremltools as ct
  
  
  class TestModel(torch.nn.Module):
      def __init__(self):
          super().__init__()
  
          self.y = torch.tensor([1, 2, 3, 4])
  
      def forward(self, x):
          idx = x.argmax()
          v = self.y[idx]
  
          return v
  
  
  if __name__ == "__main__":
  
      model = TestModel()
      x = torch.tensor([1, 2, 3, 4])
  
      print(f" Base model output: {model(x)}")
  
      traced_model = torch.jit.trace(model, x)
      out = traced_model(x)
      print(f"traced model outpus: {out}")
  
      model = ct.convert(
          traced_model,
          inputs=[ct.TensorType(name="ppg", shape=x.numpy().shape)],
      )

torch = 1.10.2 coremltools = 5.2.0

Thanks