From 6bc9858a961e1e9081883d7c48ca58107568eec7 Mon Sep 17 00:00:00 2001 From: Martin Huschenbett Date: Tue, 26 Nov 2024 01:40:44 -0800 Subject: [PATCH] Automated Code Change: Suppress pytype errors PiperOrigin-RevId: 700257299 --- examples/gemma/sampler_test.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/gemma/sampler_test.py b/examples/gemma/sampler_test.py index 2e131dda74..56870ca9bf 100644 --- a/examples/gemma/sampler_test.py +++ b/examples/gemma/sampler_test.py @@ -76,7 +76,7 @@ class SamplerTest(absltest.TestCase): def test_samples(self): vocab = MockVocab() - transformer_config = transformer_lib.TransformerConfig( + transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types num_layers=6, num_embed=vocab.GetPieceSize(), embed_dim=768, @@ -104,7 +104,7 @@ def test_samples(self): def test_forbidden_tokens(self): vocab = MockVocab() - transformer_config = transformer_lib.TransformerConfig( + transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types num_layers=0, num_embed=vocab.GetPieceSize(), embed_dim=32, @@ -152,7 +152,7 @@ def test_forbidden_tokens(self): def test_forward_equivalence(self): vocab = MockVocab() - transformer_config = transformer_lib.TransformerConfig( + transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types num_layers=2, num_embed=vocab.GetPieceSize(), embed_dim=32, @@ -211,7 +211,7 @@ def test_forward_equivalence(self): def test_sampler_init_sample_state(self): vocab = MockVocab() - transformer_config = transformer_lib.TransformerConfig( + transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types num_layers=0, num_embed=vocab.GetPieceSize(), embed_dim=32, @@ -247,7 +247,7 @@ def test_sampler_init_sample_state(self): def test_sampler_mask_tokens_after_eos_ids(self): vocab = MockVocab() - transformer_config = transformer_lib.TransformerConfig( + transformer_config = transformer_lib.TransformerConfig( # pytype: disable=wrong-arg-types num_layers=0, num_embed=vocab.GetPieceSize(), embed_dim=32,