[bnb] Fix bnb slow tests (huggingface#28788)

fix bnb slow tests
dmus · Jan 31, 2024 · d703eaa · d703eaa
1 parent 74c9cfe
commit d703eaa
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py
@@ -342,6 +342,8 @@ def to_dict(self) -> Dict[str, Any]:
         """
         output = copy.deepcopy(self.__dict__)
         output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1]
+        output["load_in_4bit"] = self.load_in_4bit
+        output["load_in_8bit"] = self.load_in_8bit
 
         return output