From 33f4cb10936f20cd92bd1e7044004c6c357fbddc Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Wed, 29 Mar 2023 15:13:00 +0200
Subject: [PATCH] [`bnb`] fix bnb failing test (#22439)

* fix bnb failing test

* fix

* fix

* fixup
---
 tests/mixed_int8/test_mixed_int8.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/mixed_int8/test_mixed_int8.py b/tests/mixed_int8/test_mixed_int8.py
index 968e2d346c..a73f689a8a 100644
--- a/tests/mixed_int8/test_mixed_int8.py
+++ b/tests/mixed_int8/test_mixed_int8.py
@@ -384,9 +384,8 @@ class MixedInt8TestMultiGpu(BaseMixedInt8Test):
         Let's just try to load a model on 2 GPUs and see if it works. The model we test has ~2GB of total, 3GB should suffice
         """
 
-        memory_mapping = {0: "1GB", 1: "2GB"}
         model_parallel = AutoModelForCausalLM.from_pretrained(
-            self.model_name, load_in_8bit=True, max_memory=memory_mapping, device_map="auto"
+            self.model_name, load_in_8bit=True, device_map="balanced"
         )
 
         # Check correct device map