IndexError: only integers, slices (`:`), ellipsis (`...`), None and long or byte Variables are valid indices (got Node)

[ueno-hiroshi002]

Issue Type

Bug

Source

pip (model-compression-toolkit)

MCT Version

2.6.0

OS Platform and Distribution

Linux Ubuntu 22.04

Python version

3.10

Describe the issue

I tried post-training quantization on a LeViT model from timm(1.0.24).

from timm import create_model
float_model = create_model("levit_conv_128s.fb_dist_in1k", pretrained=True)

However, I got the following error:

IndexError: only integers, slices (`:`), ellipsis (`...`), None and long or byte Variables are valid indices (got Node)

The following model causes the same error:

create_model("efficientformerv2_s0.snap_dist_in1k", pretrained=True)

Expected behaviour

Success PTQ.

Code to reproduce the issue

import numpy as np
from timm import create_model
from timm.data.transforms_factory import create_transform
from timm.data import resolve_data_config
import model_compression_toolkit as mct

# Prepare Model
float_model = create_model("levit_conv_128s.fb_dist_in1k", pretrained=True)
transform = create_transform(**resolve_data_config(float_model.pretrained_cfg, model=float_model), is_training=False)

# Prepare Dataset
def representative_data_gen():
    yield [np.random.random((1, 3, 224, 224))]

# MCT Quantization
quantized_model, quantization_info = mct.ptq.pytorch_post_training_quantization(
                in_module=float_model,
                representative_data_gen=representative_data_gen)

Log output

Traceback (most recent call last):
  File "/home/psnrdu/sss/model_augment/minimum_levit.py", line 16, in <module>
    quantized_model, quantization_info = mct.ptq.pytorch_post_training_quantization(
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/ptq/pytorch/quantization_facade.py", line 123, in pytorch_post_training_quantization
    tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/runner.py", line 112, in core_runner
    tg = quantization_preparation_runner(graph=graph,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/quantization_prep_runner.py", line 76, in quantization_preparation_runner
    mi.infer(_data)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/common/model_collector.py", line 232, in infer
    activation_tensors = self.fw_impl.run_model_inference(self.model, inputs_list, requires_grad=compute_hessians)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/pytorch_implementation.py", line 231, in run_model_inference
    return model(*torch_tensor_list)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py", line 377, in forward
    out_tensors_of_n, out_tensors_of_n_float = _run_operation(node,
  File "/home/psnrdu/sss/venv_mct260/py310-mct260/lib/python3.10/site-packages/model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py", line 155, in _run_operation
    out_tensors_of_n_float = op_func(*merged_inputs, **functional_kwargs)
IndexError: only integers, slices (`:`), ellipsis (`...`), None and long or byte Variables are valid indices (got Node)