[`bnb`] Let's make the daily CI green 🍏 (#21597)

* fix bnb slow test * make fixup

[`bnb`] Let's make the daily CI green 🍏 (#21597)
* fix bnb slow test * make fixup
1666c42f · Younes Belkada · GitHub · 24273268 · 1666c42f
Unverified Commit 1666c42f authored Feb 13, 2023 by Younes Belkada Committed by GitHub Feb 13, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

tests/mixed_int8/test_mixed_int8.py tests/mixed_int8/test_mixed_int8.py +3 -1

No files found.
--- a/tests/mixed_int8/test_mixed_int8.py
+++ b/tests/mixed_int8/test_mixed_int8.py
@@ -92,7 +92,9 @@ class MixedInt8Test(BaseMixedInt8Test):
        super().setUp()

        # Models and tokenizer
-        self.model_fp16 = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype="auto", device_map="auto")
+        self.model_fp16 = AutoModelForCausalLM.from_pretrained(
+            self.model_name, torch_dtype=torch.float16, device_map="auto"
+        )
        self.model_8bit = AutoModelForCausalLM.from_pretrained(self.model_name, load_in_8bit=True, device_map="auto")

    def tearDown(self):