From 6ac8e859973d9b0164cdad58038a90de0ecf545c Mon Sep 17 00:00:00 2001 From: Martin Huschenbett Date: Thu, 14 Nov 2024 05:33:45 -0800 Subject: [PATCH] Automated Code Change PiperOrigin-RevId: 696487524 --- gemma/sampler_test.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/gemma/sampler_test.py b/gemma/sampler_test.py index cadaef5..5533c2c 100644 --- a/gemma/sampler_test.py +++ b/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=1, num_embed=vocab.GetPieceSize(), embed_dim=768, @@ -112,7 +112,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=1, num_embed=vocab.GetPieceSize(), embed_dim=32, @@ -168,7 +168,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, @@ -228,7 +228,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, @@ -272,7 +272,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,