def testCheckpointLargeBatches(self): if pywrap_sanitizers.is_tsan_enabled(): self.skipTest('Creating a large buffer causes OOM when using tsan.') # Batches of size 512M dataset = dataset_ops.Dataset.from_tensors( array_ops.ones((64, 1024, 1024), dtype=dtypes.float32)).repeat() dataset = dataset.batch(2, num_parallel_calls=5) iterator = iter(dataset) next(iterator) # request an element to fill the buffer ckpt = trackable_utils.Checkpoint(iterator=iterator) manager = checkpoint_management.CheckpointManager( ckpt, self.get_temp_dir(), max_to_keep=1) manager.save()
def testDisableTsan(self): self.assertFalse(pywrap_sanitizers.is_tsan_enabled())