data_loss.py

import tensorflow as tf
from functools import partial

def collate_pair(x):
    return x[:-1], x[-1]

ds_raw = tf.constant([i for i in range(0, 100)])
ds = tf.data.Dataset.from_tensor_slices(ds_raw)

ds = ds.window(3).flat_map(lambda x: x.batch(6)).map(collate_pair)

for a, b in ds.take(2):
    print(f'a: {a} -> b: {b}')
    
# a: [0 1] -> b: 2
# a: [3 4] -> b: 5
# DATA LOSS!