Automated Code Change

PiperOrigin-RevId: 698671368
google-deepmind · Nov 21, 2024 · 06364f7 · 06364f7
1 parent 91c3739
commit 06364f7
Showing 1 changed file with 5 additions and 5 deletions.
diff --git 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,