TypeError: Negative indices with dynamic shapes on ResNet18 using ai_edge_torch

Google AI Edge
, , , ,
1

I am attempting to convert a torchvision.models.resnet18 model to TFLite using ai-edge-torch with dynamic batch size enabled via the dynamic_shapes argument.

But, it raised the error below:

TypeError: Only non-negative indices are allowed when broadcasting static shapes, but got shape (-9223372036854775808, 112, 112, 64).

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
    [... skipping hidden 14 frame]

    [... skipping hidden 13 frame]

    [... skipping hidden 3 frame]

/usr/local/lib/python3.11/dist-packages/jax/_src/lax/lax.py in _check_static_shape(shape)
     98       msg = f"Only non-negative indices are allowed when broadcasting" \
     99             f" static shapes, but got shape {shape!r}."
--> 100       raise TypeError(msg)
    101 
    102   assert shapes

TypeError: Only non-negative indices are allowed when broadcasting static shapes, but got shape (-9223372036854775808, 112, 112, 64).

While executing %sub_7 : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%clone_2, %b_bn1_running_mean), kwargs = {})
Original traceback:
  File "/usr/local/lib/python3.11/dist-packages/torchvision/models/resnet.py", line 285, in forward
    return self._forward_impl(x)
  File "/usr/local/lib/python3.11/dist-packages/torchvision/models/resnet.py", line 269, in _forward_impl
    x = self.bn1(x)

Below is the Colab environment that I used to do the conversion:

How do I solve it? Are there any steps that I did wrong? Is it possible to make the batch size dynamic?