[Gemini] more tests for Gemini (#2038)

* [Gemini] more tests for Gemini * polish code
2025-09-11 05:49:55 +00:00 · 2022-11-29 17:13:10 +08:00
parent 537e181705
commit eb7742a4bb
3 changed files with 35 additions and 26 deletions
--- a/tests/components_to_test/bert.py
+++ b/tests/components_to_test/bert.py
@@ -40,7 +40,7 @@ def get_training_components():
    num_layer = 2
    vocab_size = 32

-    def bert_model_builder(checkpoint):
+    def bert_model_builder(checkpoint: bool = False):
        config = BertConfig(vocab_size=vocab_size,
                            gradient_checkpointing=checkpoint,
                            hidden_size=hidden_dim,