def time_aranged_batch(batch_size=5):
    """Generates a batch of size 1, where the feature matrix has 64 repeated
    rows of integers arranged between 0 and 49.
    """
    device = "cuda" if torch.cuda.is_available() else "cpu"
    X = torch.stack(
        [torch.stack([torch.arange(50, device=device)] * 64)] * batch_size
    ).float()
    return X, torch.zeros(batch_size, device=device)
Generates a batch of size 1, where the feature matrix has 64 repeated
    rows of integers arranged between 0 and 49.
    
test/unit_tests/augmentation/test_transforms.py