def METHOD_NAME(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