def forward(self, local_features, ch_idxs: torch.Tensor | None = None):
        """
        Parameters
        ----------
        * local_features: (batch_size, n_chans * n_times_out, emb_dim)
        * ch_idxs: (batch_size, n_chans) | None
            Indices of the channels to use in the ``ch_names`` list passed
            as argument plus one. Index 0 is reserved for an unknown channel.

        Returns
        -------
        pos_encoding: (batch_size, n_chans * n_times_out, emb_dim)
            The first ``spat_dim`` dimensions encode the channels positional encoding
            and the following ``time_dim`` dimensions encode the temporal positional encoding.
        """
        batch_size, n_chans_times, emb_dim = local_features.shape
        if ch_idxs is None:
            ch_idxs = torch.arange(
                0,
                self.pos_encoder_spat.num_embeddings,
                device=local_features.device,
            ).repeat(batch_size, 1)

        batch_size_chs, n_chans = ch_idxs.shape
        assert emb_dim >= self.spat_dim + self.time_dim
        assert n_chans_times % n_chans == 0
        n_times = n_chans_times // n_chans

        pos_encoding = local_features.new_empty(
            (batch_size_chs, n_chans, n_times, emb_dim)
        )
        # Channel pos. encoding
        pos_encoding[:, :, :, : self.spat_dim] = self.pos_encoder_spat(ch_idxs)[
            :, :, None, :
        ]
        # Temporal pos. encoding
        self._check_encoding_time(n_times)
        _ = pos_encoding[:, :, :, self.spat_dim : self.spat_dim + self.time_dim].copy_(
            self.encoding_time[None, None, :n_times, :],
        )

        return pos_encoding.view(batch_size, n_chans_times, emb_dim)

        Parameters
        ----------
        * local_features: (batch_size, n_chans * n_times_out, emb_dim)
        * ch_idxs: (batch_size, n_chans) | None
            Indices of the channels to use in the ``ch_names`` list passed
            as argument plus one. Index 0 is reserved for an unknown channel.

        Returns
        -------
        pos_encoding: (batch_size, n_chans * n_times_out, emb_dim)
            The first ``spat_dim`` dimensions encode the channels positional encoding
            and the following ``time_dim`` dimensions encode the temporal positional encoding.
        
braindecode/models/signal_jepa.py