diff --git a/src/brevitas_examples/imagenet_classification/ptq/learned_round_utils.py b/src/brevitas_examples/imagenet_classification/ptq/learned_round_utils.py index 8494e80dc..cb5694dc3 100644 --- a/src/brevitas_examples/imagenet_classification/ptq/learned_round_utils.py +++ b/src/brevitas_examples/imagenet_classification/ptq/learned_round_utils.py @@ -30,12 +30,15 @@ import torch import torch.nn.functional as F +from brevitas import config from brevitas.core.function_wrapper.learned_round import LearnedRoundSte from brevitas.graph.calibrate import DisableEnableQuantization from brevitas.inject.enum import FloatToIntImplType from brevitas.inject.enum import LearnedRoundImplType from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer as QuantWBIOL +config.IGNORE_MISSING_KEYS = True + class StopFwdException(Exception): """Used to throw and catch an exception to stop traversing the graph."""