From 429aeb5630c48551d7c29a9a316d627f9df4b611 Mon Sep 17 00:00:00 2001
From: David Corvoysier <david@huggingface.co>
Date: Mon, 23 Dec 2024 13:03:09 +0000
Subject: [PATCH] test(decoder): add granite unit tests

---
 tests/decoder/conftest.py            | 4 ++++
 tests/decoder/test_decoder_export.py | 1 +
 2 files changed, 5 insertions(+)

diff --git a/tests/decoder/conftest.py b/tests/decoder/conftest.py
index 60d728945..677b8ffbf 100644
--- a/tests/decoder/conftest.py
+++ b/tests/decoder/conftest.py
@@ -37,6 +37,10 @@
         "model_id": "Qwen/Qwen2.5-0.5B",
         "export_kwargs": {"batch_size": 4, "sequence_length": 4096, "num_cores": 2, "auto_cast_type": "fp16"},
     },
+    "granite": {
+        "model_id": "ibm-granite/granite-3.1-2b-instruct",
+        "export_kwargs": {"batch_size": 4, "sequence_length": 4096, "num_cores": 2, "auto_cast_type": "bf16"},
+    },
     "mistral": {
         "model_id": "optimum/mistral-1.1b-testing",
         "export_kwargs": {"batch_size": 4, "sequence_length": 4096, "num_cores": 2, "auto_cast_type": "bf16"},
diff --git a/tests/decoder/test_decoder_export.py b/tests/decoder/test_decoder_export.py
index 9224ecb22..61aa57481 100644
--- a/tests/decoder/test_decoder_export.py
+++ b/tests/decoder/test_decoder_export.py
@@ -31,6 +31,7 @@
     "mixtral": "dacorvo/Mixtral-tiny",
     "opt": "hf-internal-testing/tiny-random-OPTForCausalLM",
     "qwen2": "yujiepan/qwen2.5-128k-tiny-random",
+    "granite": "hf-internal-testing/tiny-random-GraniteForCausalLM",
 }