From 26e1c702cacd97665e473e233d44fa3eb7f767b2 Mon Sep 17 00:00:00 2001 From: Benjamin Fineran Date: Mon, 16 Aug 2021 13:03:58 -0400 Subject: [PATCH] fix: initialize is_qat_recipe in preprocess_state_dict `is_qat_recipe` does not get initialized if no recipe is provided by the checkpoint --- examples/pytorch/question-answering/sparseml_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/pytorch/question-answering/sparseml_utils.py b/examples/pytorch/question-answering/sparseml_utils.py index eb32511d85f8..c9086a6c7f01 100644 --- a/examples/pytorch/question-answering/sparseml_utils.py +++ b/examples/pytorch/question-answering/sparseml_utils.py @@ -241,6 +241,8 @@ def preprocess_state_dict(pretrained_model_name_or_path): manager = ScheduledModifierManager.from_yaml(recipe) modifiers = [m.__class__.__name__ for m in manager.modifiers] is_qat_recipe = "QuantizationModifier" in modifiers + else: + is_qat_recipe = False if os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)): archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME) state_dict = torch.load(archive_file, map_location="cpu")