mct-nightly 2.1.0.20240616.434__py3-none-any.whl → 2.1.0.20240617.451__py3-none-any.whl

{mct_nightly-2.1.0.20240616.434.dist-info → mct_nightly-2.1.0.20240617.451.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240616.434
+Version: 2.1.0.20240617.451
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-2.1.0.20240616.434.dist-info → mct_nightly-2.1.0.20240617.451.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=mregiWggl4blOSBp61pEozF4QXuJnyAsxzpvrXIwh2k,1573
+model_compression_toolkit/__init__.py,sha256=1OlgfgkWMboHXhaDyyG9E0dc_vedCsy6r_gAtSq-lfY,1573
 model_compression_toolkit/constants.py,sha256=9pVleMwnhlM4QwIL2HcEq42I1uF4rlSw63RUjkxOF4w,3923
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -33,7 +33,7 @@ model_compression_toolkit/core/common/graph/__init__.py,sha256=Xr-Lt_qXMdrCnnOaU
 model_compression_toolkit/core/common/graph/base_graph.py,sha256=lmIw0srKiwCvz7KWqfwKTxyQHDy3s6rWMIXzFAa1UMo,38326
 model_compression_toolkit/core/common/graph/base_node.py,sha256=X_0zqHrKYAsmnj9tAKjVYasbFcZD8OHpjdiMj9ugQs0,29436
 model_compression_toolkit/core/common/graph/edge.py,sha256=buoSEUZwilWBK3WeBKpJ-GeDaUA1SDdOHxDpxU_bGpk,3784
-model_compression_toolkit/core/common/graph/functional_node.py,sha256=71_4TrCdqR_r0mtgxmAyqI05iP5YoQQGeSmDgynuzTw,3902
+model_compression_toolkit/core/common/graph/functional_node.py,sha256=_6HsBeLlrpLvXhLPRJswcyDa4z16-O3xzHzGuv46zBc,3897
 model_compression_toolkit/core/common/graph/graph_matchers.py,sha256=CrDoHYq4iPaflgJWmoJ1K4ziLrRogJvFTVWg8P0UcDU,4744
 model_compression_toolkit/core/common/graph/graph_searches.py,sha256=2oKuW6L8hP-oL0lFO9PhQFt9fEFgVJwpc1u4fHExAtE,5128
 model_compression_toolkit/core/common/graph/virtual_activation_weights_node.py,sha256=3el-A7j1oyoo1_9zq3faQp7IeRsFXFCvnrb3zZFXpU0,9803
@@ -186,7 +186,7 @@ model_compression_toolkit/core/keras/graph_substitutions/substitutions/weights_a
 model_compression_toolkit/core/keras/hessian/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
 model_compression_toolkit/core/keras/hessian/activation_trace_hessian_calculator_keras.py,sha256=4eJKq_Fx4mm_VuBDeeti0fTcUk1lL2yjebxCugJhvrA,8871
 model_compression_toolkit/core/keras/hessian/trace_hessian_calculator_keras.py,sha256=hRfAjgZakDaIMuERmTVjJSa_Ww6FmEudYPO9R7SuYuQ,3914
-model_compression_toolkit/core/keras/hessian/weights_trace_hessian_calculator_keras.py,sha256=KBjGr9FzyZIPD4MFtsV3LDBdJtLa0VFdIXyx_KAnjTQ,12215
+model_compression_toolkit/core/keras/hessian/weights_trace_hessian_calculator_keras.py,sha256=Xogd90kPZPvKbplZQv5B77Dq_m4aW5-bL6Jxh33VZWs,12213
 model_compression_toolkit/core/keras/mixed_precision/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
 model_compression_toolkit/core/keras/mixed_precision/configurable_activation_quantizer.py,sha256=aW8wR13fK6P6xzbU9XGU60IO1yYzXSo_Hk4qeq486kg,5137
 model_compression_toolkit/core/keras/mixed_precision/configurable_weights_quantizer.py,sha256=Ziydik2j-LvNBXP3TSfUD6rEezPAikzQGib0_IXkmGM,6729
@@ -210,7 +210,7 @@ model_compression_toolkit/core/keras/statistics_correction/__init__.py,sha256=9H
 model_compression_toolkit/core/keras/statistics_correction/apply_second_moment_correction.py,sha256=XNCtT9klMcsO1v5KA3MmCq_WgXOIT5QSzbfTOa9T-04,3060
 model_compression_toolkit/core/keras/visualization/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/pytorch/constants.py,sha256=NI-J7REuxn06oEIHsmJ4GqtNC3TbV8xlkJjt5Ar-c4U,2626
+model_compression_toolkit/core/pytorch/constants.py,sha256=AguUnAsNlj41gwuKIP_7nos3FcJHsIAjewLXSQdrDQM,2624
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=r1XyzUFvrjGcJHQM5ETLsMZIG2yHCr9HMjqf0ti9inw,4175
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=S25cuw10AW3SEN_fRAGRcG_I3wdvvQx1ehSJzPnn-UI,4404
 model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=7CFt1Y3fiDaKkEVvlDd76ZmucCuVp6OZNQwwqJezKbU,27547
@@ -222,7 +222,7 @@ model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py,s
 model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py,sha256=tLrlUyYhxVKVjkad1ZAtbRra0HedB3iVfIkZ_dYnQ-4,3419
 model_compression_toolkit/core/pytorch/back2framework/instance_builder.py,sha256=BBHBfTqeWm7L3iDyPBpk0jxvj-rBg1QWI23imkjfIl0,1467
 model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py,sha256=D7lU1r9Uq_7fdNuKk2BMF8ho5GrsY-8gyGN6yYoHaVg,15060
-model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=Zw4gi-wjJNV8-qGv79YBWVAHmy27f7iW0c2JGNWAKD0,18199
+model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py,sha256=Pvpkirt3ziWEXDEspgOhR8ALf-XAZUh-78IkXg9YMWs,18830
 model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py,sha256=qZNNOlNTTV4ZKPG3q5GDXkIVTPUEr8dvxAS_YiMORmg,3456
 model_compression_toolkit/core/pytorch/back2framework/quantization_wrapper/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/core/pytorch/back2framework/quantization_wrapper/quantized_layer_wrapper.py,sha256=q2JDw10NKng50ee2i9faGzWZ-IydnR2aOMGSn9RoZmc,5773
@@ -240,7 +240,7 @@ model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py,sha256=VNg-VgzCxSyqy2J3neEPl6U0SPO8UIVU_T47bGhz4FE,38459
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py,sha256=q1a3HieQtaOmWG2WGXp6GHYAvxa3CZ9dJUx9dqMAsS8,5695
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/remove_identity.py,sha256=joHjwiUxccypMHkTy46rI91VyapLn9yJ2YRo5ISnOH4,1987
-model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/reshape_with_static_shapes.py,sha256=jOqlelGhADEZiYUEyYj9oJZ5YLXx8jWNUlVTG6Td79Y,4919
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/reshape_with_static_shapes.py,sha256=hAZXzrEinHa-dJHLj39Hy_9Q-13QyO95rtYVSLrhvT8,4915
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/residual_collapsing.py,sha256=DcJEIkGvBdIMOelNIwaJUZ5UsAHiGnDJPR20I464vWo,2929
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/scale_equalization.py,sha256=XFtU9yuBmoZlX0f0mS6otMPWMk-RcWs94XdvvTNhW8Y,3303
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/shift_negative_activation.py,sha256=lOPl5zDU3FoR9WmlxO04Pfi65MimK0gmnuHzQJodQdY,10668
@@ -261,7 +261,7 @@ model_compression_toolkit/core/pytorch/quantizer/__init__.py,sha256=Rf1RcYmelmdZ
 model_compression_toolkit/core/pytorch/quantizer/fake_quant_builder.py,sha256=D8_CEuFqKAhbUgKaRw7Jlxo0zlqgPTMu6CIIIM4LfS0,7045
 model_compression_toolkit/core/pytorch/quantizer/lut_fake_quant.py,sha256=uyeBtNokyDUikk-YkDP_mN_2DX0J5oPm3kSfdSUT2Ck,4420
 model_compression_toolkit/core/pytorch/reader/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
-model_compression_toolkit/core/pytorch/reader/graph_builders.py,sha256=x5n8KHBqvutqS5l5AillA_FQfhf-2ibP813ixK3Gvy8,12627
+model_compression_toolkit/core/pytorch/reader/graph_builders.py,sha256=LiGV-ZqlhxN1evpM-ur2dDVPowhrLwO7JZa7AGPftSk,12913
 model_compression_toolkit/core/pytorch/reader/node_holders.py,sha256=TaolORuwBZEddWe-q0Mg79Nmswz-Sq3-9-4o8UxFQ50,1028
 model_compression_toolkit/core/pytorch/reader/reader.py,sha256=GEJE0QX8XJFWbYCkbRBtzttZtmmuoACLx8gw9KyAQCE,6015
 model_compression_toolkit/core/pytorch/statistics_correction/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
@@ -491,8 +491,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=MxylaVFPgN7zBiRBy6WV610EA4scLgRJFbMucKvvNDU,2896
-mct_nightly-2.1.0.20240616.434.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.1.0.20240616.434.dist-info/METADATA,sha256=qRPmiufe7bR0fBvN5cB16-cYIewTIbYTbbpfBQvjA_8,19721
-mct_nightly-2.1.0.20240616.434.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.1.0.20240616.434.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.1.0.20240616.434.dist-info/RECORD,,
+mct_nightly-2.1.0.20240617.451.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.1.0.20240617.451.dist-info/METADATA,sha256=-ZUI2y7SZOGKyLl6qpBE9onj-lZwfGE0wLLJI5WeqIE,19721
+mct_nightly-2.1.0.20240617.451.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.1.0.20240617.451.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.1.0.20240617.451.dist-info/RECORD,,

model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.1.0.20240616.000434"
+__version__ = "2.1.0.20240617.000451"

model_compression_toolkit/core/common/graph/functional_node.py

@@ -25,7 +25,7 @@ class FunctionalNode(BaseNode):
                  functional_op: Any = None,
                  inputs_as_list: bool = False,
                  has_activation: bool = True,
-                 tensor_input_indices = None):
+                 tensor_input_allocs = None):
         """
         Init a FunctionalNode object.
 
@@ -44,7 +44,7 @@ class FunctionalNode(BaseNode):
             functional_op: The op the node implements.
             inputs_as_list: Whether to pass the node its input tensors as a list or not when calling the layer.
             has_activation: Whether the node has activations that we might want to quantize.
-            tensor_input_indices: A list of indices for activation tensors in the node's input tensor list
+            tensor_input_allocs: A list of indices for activation tensors in the node's input tensor list
 
         """
 
@@ -63,7 +63,7 @@ class FunctionalNode(BaseNode):
         self.op_call_args = op_call_args
         self.functional_op = functional_op
         self.inputs_as_list = inputs_as_list
-        self.tensor_input_indices = [] if tensor_input_indices is None else tensor_input_indices
+        self.tensor_input_allocs = [] if tensor_input_allocs is None else tensor_input_allocs
 
     @property
     def type(self):

model_compression_toolkit/core/keras/hessian/weights_trace_hessian_calculator_keras.py

@@ -74,7 +74,7 @@ class WeightsTraceHessianCalculatorKeras(TraceHessianCalculatorKeras):
         model, _ = FloatKerasModelBuilder(graph=self.graph).build_model()
 
         # Initiate a gradient tape for automatic differentiation
-        with (tf.GradientTape(persistent=True) as tape):
+        with tf.GradientTape(persistent=True) as tape:
             # Perform a forward pass (inference) to get the output, while watching
             # the input tensor for gradient computation
             tape.watch(self.input_images)

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -22,6 +22,7 @@ from networkx import topological_sort
 
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.back2framework.base_model_builder import BaseModelBuilder
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
@@ -66,8 +67,8 @@ def _build_input_tensors_list(node: BaseNode,
     return input_tensors
 
 
-def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List,
-                  tensor_input_indices: List = None) -> List:
+def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List, op_call_kwargs: Dict,
+                  tensor_input_allocs: List = None) -> Tuple[List, Dict]:
     """
     Merge input tensors list with positional weights and op_call_args, according to correct order.
 
@@ -75,22 +76,30 @@ def _merge_inputs(_node: BaseNode, input_tensors: List, op_call_args: List,
         _node: The node the inputs are for.
         input_tensors: activation input tensors to node.
         op_call_args: framework node call args.
+        op_call_kwargs: framework node call kwargs.
+        tensor_input_allocs: List of input allocations to node.
 
     Returns:
         Combined list of input_tensors and op_call_args.
     """
-    if isinstance(_node, FunctionalNode) and _node.tensor_input_indices:
+    if isinstance(_node, FunctionalNode) and _node.tensor_input_allocs:
         _input_list = op_call_args.copy()
-        if tensor_input_indices is None:
-            tensor_input_indices = _node.tensor_input_indices
-        assert len(tensor_input_indices) == len(input_tensors), \
-            f'Mismatch between input tensors ({len(tensor_input_indices)}) and indices {len(input_tensors)}'
-        for i, t in zip(tensor_input_indices, input_tensors):
-            _input_list.insert(i, t)
+        if tensor_input_allocs is None:
+            tensor_input_allocs = _node.tensor_input_allocs
+        if len(tensor_input_allocs) != len(input_tensors):
+            Logger.error(f'Mismatch between input tensors ({len(tensor_input_allocs)}) '
+                         f'and indices {len(input_tensors)} in node {_node.name}.')  # pragma: no cover
+        for i, t in zip(tensor_input_allocs, input_tensors):
+            # insert input tensors in either args or kwargs, according to tensor_input_allocs
+            if isinstance(i, str):
+                assert i not in op_call_kwargs
+                op_call_kwargs.update({i: t})
+            else:
+                _input_list.insert(i, t)
     else:
         _input_list = input_tensors + op_call_args
 
-    return _input_list
+    return _input_list, op_call_kwargs
 
 
 def _run_operation(n: BaseNode,
@@ -125,14 +134,15 @@ def _run_operation(n: BaseNode,
         # list separately, because in FX the tensors are FX objects and fail to_torch_tensor
         input_tensors = [to_torch_tensor(t, numpy_type=t.dtype) if isinstance(t, np.ndarray) else t
                          for t in input_tensors]
-        _tensor_input_indices = None
+        _tensor_input_allocs = None
     else:
-        _tensor_input_indices = [i for i in n.tensor_input_indices if i not in n.weights]
+        _tensor_input_allocs = [i for i in n.tensor_input_allocs if i not in n.weights]
 
     if isinstance(n, FunctionalNode) and n.inputs_as_list:
         out_tensors_of_n_float = op_func(input_tensors, *op_call_args, **functional_kwargs)
     else:
-        merged_inputs = _merge_inputs(n, input_tensors, op_call_args, tensor_input_indices=_tensor_input_indices)
+        merged_inputs, functional_kwargs = _merge_inputs(n, input_tensors, op_call_args, functional_kwargs.copy(),
+                                                         tensor_input_allocs=_tensor_input_allocs)
         out_tensors_of_n_float = op_func(*merged_inputs, **functional_kwargs)
 
     # Add a fake quant node if the node has an activation threshold.

model_compression_toolkit/core/pytorch/constants.py

@@ -40,7 +40,7 @@ FUNCTIONAL_OP = 'functional_op'
 OP_CALL_ARGS = 'op_call_args'
 OP_CALL_KWARGS = 'op_call_kwargs'
 INPUTS_AS_LIST = 'inputs_as_list'
-TENSOR_INPUT_INDICES = 'tensor_input_indices'
+TENSOR_INPUT_ALLOCS = 'tensor_input_allocs'
 INPLACE = 'inplace'
 HARDTANH_MIN_VAL = 'min_val'
 HARDTANH_MAX_VAL = 'max_val'

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/reshape_with_static_shapes.py

@@ -65,10 +65,10 @@ class ReshapeWithStaticShapes(common.BaseSubstitution):
 
         # When a "reshape" is called with multiple arguments (e.g. x.reshape(-1, channels, height, width)
         # this substitution converts it x.reshape((-1, channels, height, width)), so need to update the
-        # tensor_input_indices attribute.
-        # scalar argument's shape is [1] so remove those indices from tensor_input_indices
+        # tensor_input_allocs attribute.
+        # scalar argument's shape is [1] so remove those indices from tensor_input_allocs
         # node.input_shape example: [[1, 32, 4, 32], [1], [1], [1]]
-        node.tensor_input_indices = node.tensor_input_indices[:sum([i != [1] for i in node.input_shape])]
+        node.tensor_input_allocs = node.tensor_input_allocs[:sum([i != [1] for i in node.input_shape])]
 
         # modify the node input info
         node.input_shape = [node.input_shape[0]]

model_compression_toolkit/core/pytorch/reader/graph_builders.py

@@ -23,7 +23,7 @@ from model_compression_toolkit.core.common.graph.base_graph import OutTensor
 from model_compression_toolkit.core.common.graph.edge import Edge
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.constants import OUTPUT, PLACEHOLDER, TENSOR_META, CALL_FUNCTION, TYPE, \
-    CALL_METHOD, BIAS, FUNCTIONAL_OP, OP_CALL_KWARGS, OP_CALL_ARGS, INPUTS_AS_LIST, TENSOR_INPUT_INDICES, GET_ATTR
+    CALL_METHOD, BIAS, FUNCTIONAL_OP, OP_CALL_KWARGS, OP_CALL_ARGS, INPUTS_AS_LIST, TENSOR_INPUT_ALLOCS, GET_ATTR
 from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder
 from model_compression_toolkit.logger import Logger
 
@@ -140,7 +140,7 @@ def nodes_builder(model: GraphModule,
                     weights.update({i: consts_dict[input_node]})
 
                 tensor_meta = input_node.meta
-                if tensor_meta[TYPE] == torch.Tensor:
+                if tensor_meta[TYPE] in [torch.Tensor, torch.nn.parameter.Parameter]:
                     input_shape += [list(tensor_meta[TENSOR_META].shape)]
                 elif tensor_meta[TYPE] == tuple:
                     input_shape += [list(n.shape) for n in tensor_meta[TENSOR_META]]
@@ -159,8 +159,11 @@ def nodes_builder(model: GraphModule,
 
         # filter Nodes from framework attributes, we replace these attributes with nx graph nodes
         framework_attr_filtered = {}
+        framework_attr_nodes = {}
         for k, v in framework_attr.items():
-            if not isinstance(v, torch.fx.node.Node):
+            if isinstance(v, torch.fx.node.Node):
+                framework_attr_nodes[k] = v
+            else:
                 framework_attr_filtered[k] = v
         framework_attr = framework_attr_filtered
 
@@ -177,7 +180,7 @@ def nodes_builder(model: GraphModule,
                 [isinstance(n, torch.fx.node.Node) for n in node.args[0]])
             inputs_as_list = inputs_as_list1 or (len(node.args) > 0 and isinstance(node.args[0], Node) and
                                                  node.args[0].op == PLACEHOLDER and node.args[0].meta[TYPE] in (list, tuple))
-            tensor_input_index = []
+            tensor_input_alloc = []
             op_call_args = list(node.args)
             if inputs_as_list:
                 op_call_args.pop(0)
@@ -185,7 +188,10 @@ def nodes_builder(model: GraphModule,
                 for in_node in node.all_input_nodes:
                     for i, arg in enumerate(node.args):
                         if arg == in_node:
-                            tensor_input_index.append(i)
+                            tensor_input_alloc.append(i)
+                    for k, arg in framework_attr_nodes.items():
+                        if arg == in_node:
+                            tensor_input_alloc.append(k)
 
             # remove torch.fx.node.Node from inputs to graph_node_type
             op_call_args = [arg for arg in op_call_args if not isinstance(arg, Node)]
@@ -197,7 +203,7 @@ def nodes_builder(model: GraphModule,
                       OP_CALL_ARGS: op_call_args,
                       OP_CALL_KWARGS: node_kwargs,
                       INPUTS_AS_LIST: inputs_as_list,
-                      TENSOR_INPUT_INDICES: tensor_input_index}
+                      TENSOR_INPUT_ALLOCS: tensor_input_alloc}
         else:
             graph_node_type = BaseNode
             kwargs = {}