...cleanup and resolve pytype issues with assertions...

PiperOrigin-RevId: 689850257
google-research · Nov 21, 2024 · 0605ba6 · 0605ba6
1 parent 0340172
commit 0605ba6
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/scenic/dataset_lib/tests/test_dataset_utils.py b/scenic/dataset_lib/tests/test_dataset_utils.py
@@ -74,6 +74,7 @@ def make_fake_batches():
     self.assertEqual(outputs['batch_mask'].shape,
                      (desired_bs,) + batch_mask_shape)
     if pre_padding_mask:
+      assert complete_batch_mask is not None
       self.assertEqual(outputs['batch_mask'].sum(), complete_batch_mask.sum())
     else:
       self.assertEqual(outputs['batch_mask'].sum(),
@@ -104,6 +105,7 @@ def make_fake_batches():
     self.assertEqual(outputs['batch_mask'].shape,
                      (desired_bs,) + batch_mask_shape)
     if pre_padding_mask:
+      assert complete_batch_mask is not None
       self.assertEqual(outputs['batch_mask'].sum(), complete_batch_mask.sum())
     else:
       self.assertEqual(outputs['batch_mask'].sum(),