From 7f4e69395c1afab759d1e16aa62f868004cb1a9d Mon Sep 17 00:00:00 2001 From: Ashvin Kumar Date: Wed, 25 Oct 2023 11:24:37 -0700 Subject: [PATCH] Removed typehints to support upgraded onnx version Signed-off-by: Ashvin Kumar --- .../common/src/python/aimet_common/cross_layer_equalization.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/TrainingExtensions/common/src/python/aimet_common/cross_layer_equalization.py b/TrainingExtensions/common/src/python/aimet_common/cross_layer_equalization.py index c681a90d4ea..e92227526e2 100644 --- a/TrainingExtensions/common/src/python/aimet_common/cross_layer_equalization.py +++ b/TrainingExtensions/common/src/python/aimet_common/cross_layer_equalization.py @@ -46,7 +46,6 @@ from typing import List, Union, Tuple, Dict from enum import Enum import numpy as np -from onnx import onnx_pb from aimet_common.utils import AimetLogger from aimet_common.connected_graph.connectedgraph_utils import get_all_input_ops @@ -76,7 +75,7 @@ class ClsSetLayerPairInfo: Models a pair of layers that were scaled using CLS. And related information. """ - def __init__(self, layer1: onnx_pb.NodeProto, layer2: onnx_pb.NodeProto, scale_factor: np.ndarray, + def __init__(self, layer1, layer2, scale_factor: np.ndarray, relu_activation_between_layers: bool): """ :param layer1: Layer whose bias is folded