From 5ce90f32122f98942e3b54f0c4fbc3b393b396aa Mon Sep 17 00:00:00 2001
From: Titus <9048635+Titus-von-Koeller@users.noreply.github.com>
Date: Mon, 19 Feb 2024 15:04:44 -0300
Subject: [PATCH] Bnb test fix for different hardwares (#29066)

* generated text on A10G

* generated text in CI

* Apply suggestions from code review

add explanatory comments

Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>

---------

Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
---
 tests/quantization/bnb/test_mixed_int8.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py
index 1c4362ef0c..cf6e156034 100644
--- a/tests/quantization/bnb/test_mixed_int8.py
+++ b/tests/quantization/bnb/test_mixed_int8.py
@@ -97,6 +97,8 @@ class BaseMixedInt8Test(unittest.TestCase):
     input_text = "Hello my name is"
     EXPECTED_OUTPUTS = set()
     EXPECTED_OUTPUTS.add("Hello my name is John.\nI am a friend of the family.\n")
+    # Expected values on a A10
+    EXPECTED_OUTPUTS.add("Hello my name is John.\nI am a friend of your father.\n")
     MAX_NEW_TOKENS = 10
 
     def setUp(self):
@@ -847,6 +849,8 @@ class MixedInt8GPT2Test(MixedInt8Test):
     EXPECTED_OUTPUTS = set()
     EXPECTED_OUTPUTS.add("Hello my name is John Doe, and I'm a big fan of")
     EXPECTED_OUTPUTS.add("Hello my name is John Doe, and I'm a fan of the")
+    # Expected values on a A10
+    EXPECTED_OUTPUTS.add("Hello my name is John Doe, and I am a member of the")
 
     def test_int8_from_pretrained(self):
         r"""