import torch
from torch.utils.data import DataLoader, TensorDataset


features = torch.arange(24, dtype=torch.float32).reshape(6, 4)
labels = torch.arange(6)
dataset = TensorDataset(features, labels)

loader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    generator=torch.Generator().manual_seed(17),
)

batch_features, batch_labels = next(iter(loader))

print(f"dataset length: {len(dataset)}")
print(f"batch features shape: {tuple(batch_features.shape)}")
print(f"batch labels: {batch_labels.tolist()}")
print("batch features:")
print(batch_features)

for epoch in range(2):
    epoch_loader = DataLoader(
        dataset,
        batch_size=2,
        shuffle=True,
        generator=torch.Generator().manual_seed(17),
    )
    label_order = []
    for _, batch_targets in epoch_loader:
        label_order.extend(batch_targets.tolist())
    print(f"epoch {epoch} label order: {label_order}")