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