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

PiperOrigin-RevId: 689850257

Author: Scenic Authors

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(),