Converting model to nnapi, failing at serializing

I’m following (Beta) Convert MobileNetV2 to NNAPI — PyTorch Tutorials 2.0.1+cu117 documentation to get my model to work on android gpu device but I’m using

        traced = torch.jit.script(model, input_tensor)

Instead of trace, as it’s supposed to be more suitable for my model architecture (YoloX), I’ve removed the head since it’s supposed to be iffy to add it too, and converted depthwise convs to regular 2dconve both for support and as it’s supposed to run better on gpu operations, but I’m getting the following error when running the following line:

    nnapi_model = torch.backends._nnapi.prepare.convert_model_to_nnapi(traced, input_tensor)

  retval_count) = serializer.serialize_model(model, inputs, return_shapes)
  File "/envs/yolox/lib/python3.9/site-packages/torch/backends/_nnapi/serializer.py", line 677, in serialize_model
    self.add_node(node)
  File "/envs/yolox/lib/python3.9/site-packages/torch/backends/_nnapi/serializer.py", line 871, in add_node
    adder(self, node)
  File "/envs/yolox/lib/python3.9/site-packages/torch/backends/_nnapi/serializer.py", line 848, in <lambda>
    self.add_conv2d(node),
  File "/envs/yolox/lib/python3.9/site-packages/torch/backends/_nnapi/serializer.py", line 1779, in add_conv2d
    return self.add_conv2d_common(
  File "/envs/yolox/lib/python3.9/site-packages/torch/backends/_nnapi/serializer.py", line 1979, in add_conv2d_common
    assert kern_d == in_c
AssertionError

I’ve added a print to those, and I’m getting kern_d=32, in_c=64

Any ideas how to correct this issue?
Thanks.