def forward(self, X: Tensor, y: Optional[Tensor] = None) -> Output:
        """General forward pass for an augmentation transform.

        Parameters
        ----------
        X : torch.Tensor
            EEG input example or batch.
        y : torch.Tensor | None
            EEG labels for the example or batch. Defaults to None.

        Returns
        -------
        torch.Tensor
            Transformed inputs.
        torch.Tensor, optional
            Transformed labels. Only returned when y is not None.
        """
        X = torch.as_tensor(X).float()

        out_X = X.clone()
        # check if input has a batch dimension
        if len(out_X.shape) < 3:
            out_X = out_X[None, ...]

        if y is not None:
            y = torch.as_tensor(y).to(out_X.device)
            out_y = y.clone()
            if len(out_y.shape) == 0:
                out_y = out_y.reshape(1)
        else:
            out_y = torch.zeros(out_X.shape[0], device=out_X.device)

        # Samples a mask setting for each example whether they should stay
        # unchanged or not
        mask = self._get_mask(out_X.shape[0], out_X.device)
        num_valid = mask.sum().long()

        if num_valid > 0:
            # Uses the mask to define the output
            out_X[mask, ...], tr_y = self.operation(
                out_X[mask, ...],
                out_y[mask],
                **self.get_augmentation_params(out_X[mask, ...], out_y[mask]),
            )
            # Apply the operation defining the Transform to the whole batch
            if isinstance(tr_y, tuple):
                out_y = tuple(tmp_y[mask] for tmp_y in tr_y)
            else:
                out_y[mask] = tr_y

        # potentially remove empty batch dimension again
        out_X = out_X.reshape_as(X)
        if y is not None:
            return out_X, out_y
        else:
            return out_X
General forward pass for an augmentation transform.

        Parameters
        ----------
        X : torch.Tensor
            EEG input example or batch.
        y : torch.Tensor | None
            EEG labels for the example or batch. Defaults to None.

        Returns
        -------
        torch.Tensor
            Transformed inputs.
        torch.Tensor, optional
            Transformed labels. Only returned when y is not None.
        
braindecode/augmentation/base.py