CustomMarker module with list as output

[lkdci]

aimet-torch version: AimetTorch-torch-gpu-1.29.0 env: torch-gpu docker environment

I tried to call the QuantizationSimModel.export() method for the YoloNAS model, unfortunately there is a sub module in yolonas which output a list of tensor, see this line

This cause the CustomMarker to return a wrong output when the output list contains only one item, as only output type of tuple are handled for this use case, as in the following aimet code:

was_output_tuple = isinstance(x, tuple)
...
if len(output) == 1 and not was_output_tuple:
    output = output[0]

I believe this can be easily solved by checking also for a list type, or alternatively for an Iterable type:

was_output_tuple = isinstance(x, (tuple, list))

Unfortunately I couldn't find a way to build the aimet-torch for gpu from source, so I can test and resolve this bug, until it can be merged into a release tag.

Any idea how I can work around it?

[quic-hitameht]

Hi @lkdci,

You can follow instructions provided in this documentation to build and install latest AIMET source code. Once, you have a PR ready with the fix, please let us know.

Thank you.