def random_batch(rng_seed, batch_size=5):
    """Generate batch of elements containing feature matrix of size 66x50
    filled with random floats between 0 and 1.
    """
    device = "cuda" if torch.cuda.is_available() else "cpu"
    rng = check_random_state(rng_seed)
    X = torch.from_numpy(rng.random((batch_size, 22, 51))).float().to(device)
    return X, torch.zeros(batch_size, device=device)
Generate batch of elements containing feature matrix of size 66x50
    filled with random floats between 0 and 1.
    
test/unit_tests/augmentation/conftest.py