No public description

PiperOrigin-RevId: 699280431
google-research · Feb 27, 2025 · 39bebf2 · 39bebf2
1 parent 47c4738
commit 39bebf2
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:
+      self.assertIsNotNone(complete_batch_mask)
       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:
+      self.assertIsNotNone(complete_batch_mask)
       self.assertEqual(outputs['batch_mask'].sum(), complete_batch_mask.sum())
     else:
       self.assertEqual(outputs['batch_mask'].sum(),