From 90b20fc1dbd056c999501e5363fec90eb12f2315 Mon Sep 17 00:00:00 2001 From: Pablo Montalvo Date: Fri, 9 Aug 2024 11:13:17 +0200 Subject: [PATCH 1/2] add Gemma2 for causal LM to test classes --- tests/models/gemma2/test_modeling_gemma2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/gemma2/test_modeling_gemma2.py b/tests/models/gemma2/test_modeling_gemma2.py index 1229ca47eb69..03e4000fc21f 100644 --- a/tests/models/gemma2/test_modeling_gemma2.py +++ b/tests/models/gemma2/test_modeling_gemma2.py @@ -59,7 +59,7 @@ class Gemma2ModelTest(GemmaModelTest, unittest.TestCase): if is_torch_available() else () ) - all_generative_model_classes = () + all_generative_model_classes = (Gemma2ForCausalLM,) if is_torch_available else () pipeline_model_mapping = ( { "feature-extraction": Gemma2Model, From 5b64dfe71868fc98a8c291f25473053ba760616d Mon Sep 17 00:00:00 2001 From: Pablo Montalvo Date: Fri, 9 Aug 2024 11:20:38 +0200 Subject: [PATCH 2/2] fix --- tests/models/gemma2/test_modeling_gemma2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/gemma2/test_modeling_gemma2.py b/tests/models/gemma2/test_modeling_gemma2.py index 03e4000fc21f..71c61bff8f24 100644 --- a/tests/models/gemma2/test_modeling_gemma2.py +++ b/tests/models/gemma2/test_modeling_gemma2.py @@ -59,7 +59,7 @@ class Gemma2ModelTest(GemmaModelTest, unittest.TestCase): if is_torch_available() else () ) - all_generative_model_classes = (Gemma2ForCausalLM,) if is_torch_available else () + all_generative_model_classes = (Gemma2ForCausalLM,) if is_torch_available() else () pipeline_model_mapping = ( { "feature-extraction": Gemma2Model,