Automated Code Change: Suppress pytype errors

PiperOrigin-RevId: 700257299
google · Nov 26, 2024 · 6bc9858 · 6bc9858
1 parent abc1155
commit 6bc9858
Showing 1 changed file with 5 additions and 5 deletions.
diff --git 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,