def _make_collateable(transform, device=None):
    """Wraps a transform to make it collateable.
    with device control."""

    def _collate_fn(batch):
        collated_batch = default_collate(batch)
        X, y = collated_batch[:2]

        if device is not None:
            X = X.to(device)
            y = y.to(device)

        return (*transform(X, y), *collated_batch[2:])

    return _collate_fn
Wraps a transform to make it collateable.
    with device control.
braindecode/augmentation/base.py