FIX [bnb / tests]: Fix currently failing bnb tests (huggingface#2…

…9092) Update test_mixed_int8.py
zucchini-nlp · Feb 19, 2024 · 457945f · 457945f
1 parent 1df37a1
commit 457945f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py
@@ -42,7 +42,7 @@
 
 
 def get_some_linear_layer(model):
-    if model.config.model_type == "openai-community/gpt2":
+    if model.config.model_type == "gpt2":
         return model.transformer.h[0].mlp.c_fc
     return model.transformer.h[0].mlp.dense_4h_to_h