Braindecode 1.7.0__tar.gz → 1.7.0.dev174519142__tar.gz

{braindecode-1.7.0/braindecode.egg-info → braindecode-1.7.0.dev174519142}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.7.0
+Version: 1.7.0.dev174519142
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/classifier.py

@@ -3,7 +3,6 @@
 #          Lukas Gemein <l.gemein@gmail.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
 #          Pierre Guetschel <pierre.guetschel@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -196,7 +195,9 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
         return self.predict_proba(X).argmax(1)
 
     def predict_trials(self, X, return_targets=True):
-        """Create trialwise predictions from a cropped dataset.
+        """Create trialwise predictions and optionally also return trialwise.
+
+        labels from cropped dataset.
 
         Parameters
         ----------
@@ -211,11 +212,10 @@ class EEGClassifier(_EEGNeuralNet, NeuralNetClassifier):
             3-dimensional array (n_trials x n_classes x n_predictions), where
             the number of predictions depend on the chosen window size and the
             receptive field of the network.
-        trial_targets : np.ndarray
-            Ground-truth targets from the dataset in a 2-dimensional array
-            (n_trials x n_targets). Only returned when ``return_targets=True``.
-            The number of targets depends on the decoding paradigm and can be
-            either a single value, multiple values, or a sequence.
+        trial_labels : np.ndarray
+            2-dimensional array (n_trials x n_targets) where the number of
+            targets depends on the decoding paradigm and can be either a single
+            value, multiple values, or a sequence.
         """
         if not self.cropped:
             warnings.warn(

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/functional/__init__.py

@@ -5,8 +5,6 @@ from .functions import (
     dwt_max_level,
     hilbert_freq,
     identity,
-    iou_1d,
-    pairwise_iou_1d,
     plv_time,
     safe_log,
     sinusoidal_positional_encoding,
@@ -22,8 +20,6 @@ __all__ = [
     "dwt_max_level",
     "hilbert_freq",
     "identity",
-    "iou_1d",
-    "pairwise_iou_1d",
     "plv_time",
     "safe_log",
     "sinusoidal_positional_encoding",

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/functional/functions.py

@@ -1,6 +1,5 @@
 # Authors: Robin Schirrmeister <robintibor@gmail.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -110,7 +109,7 @@ def _get_gaussian_kernel1d(kernel_size: int, sigma: float) -> torch.Tensor:
     return kernel1d
 
 
-def hilbert_freq(x: torch.Tensor, forward_fourier: bool = True) -> torch.Tensor:
+def hilbert_freq(x, forward_fourier=True):
     r"""
     Compute the Hilbert transform using PyTorch, separating the real and
     imaginary parts.
@@ -179,9 +178,7 @@ def hilbert_freq(x: torch.Tensor, forward_fourier: bool = True) -> torch.Tensor:
     return x
 
 
-def plv_time(
-    x: torch.Tensor, forward_fourier: bool = True, epsilon: float = 1e-6
-) -> torch.Tensor:
+def plv_time(x, forward_fourier=True, epsilon: float = 1e-6):
     """Compute the Phase Locking Value (PLV) metric in the time domain.
 
     The Phase Locking Value (PLV) is a measure of the synchronization between
@@ -260,40 +257,6 @@ def plv_time(
     return plv_matrix
 
 
-# -----------------------------------------------------------------------------
-# DANCE functional helpers
-#
-# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
-#
-# License: MIT
-#
-# Ported from the DANCE event-detection model (facebookresearch/dance, MIT).
-# -----------------------------------------------------------------------------
-
-
-def iou_1d(s1, e1, s2, e2, eps: float = 1e-7):
-    """ELEMENTWISE 1-D temporal IoU. All four inputs share the same shape S
-    (e.g. ``(B, Q)``); returns IoU of shape S. Used by ``DanceLoss`` on the
-    matched ``(B, Q)`` spans. For (Q,)x(T,) pairwise IoU use ``pairwise_iou_1d``.
-    """
-    inter = (torch.minimum(e1, e2) - torch.maximum(s1, s2)).clamp(min=0)
-    union = (e1 - s1) + (e2 - s2) - inter
-    return inter / (union + eps)
-
-
-def pairwise_iou_1d(s1, e1, s2, e2, eps: float = 1e-7):
-    """PAIRWISE 1-D temporal IoU. ``s1,e1`` shape ``(Q,)``, ``s2,e2`` shape
-    ``(T,)``; returns ``(Q, T)`` (broadcast ``[:, None]`` x ``[None, :]``).
-    Transcribed from ``dance/matcher.py:33-39`` (``_pairwise_iou``). Used by
-    ``HungarianMatcher`` to build the ``(Q, n_targets)`` cost matrix.
-    """
-    inter_start = torch.maximum(s1[:, None], s2[None, :])
-    inter_end = torch.minimum(e1[:, None], e2[None, :])
-    inter = (inter_end - inter_start).clamp(min=0)
-    union = (e1 - s1)[:, None] + (e2 - s2)[None, :] - inter
-    return inter / (union + eps)
-
-
 def daubechies_filters(n_vanishing: int) -> torch.Tensor:
     r"""Daubechies ``db<n_vanishing>`` wavelet decomposition filters.
 

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/models/__init__.py

@@ -11,7 +11,6 @@ from .cbramod import CBraMod
 from .codebrain import CodeBrain
 from .contrawr import ContraWR
 from .ctnet import CTNet
-from .dance import DANCE
 from .deep4 import Deep4Net
 from .deepsleepnet import DeepSleepNet
 from .dgcnn import DGCNN
@@ -87,7 +86,6 @@ __all__ = [
     "CodeBrain",
     "ContraWR",
     "CTNet",
-    "DANCE",
     "Deep4Net",
     "DeepSleepNet",
     "BrainModule",

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/models/bendr.py

@@ -1,8 +1,3 @@
-# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
-#
-# License: BSD (3-clause)
-
 import copy
 
 import numpy as np
@@ -440,9 +435,7 @@ class _ConvEncoderBENDR(nn.Module):
                         padding=width
                         // 2,  # Correct padding for 'same' output length before stride
                     ),
-                    # channel-wise dropout (matches paper specification), on
-                    # (batch, channels, times) activations
-                    nn.Dropout1d(dropout),
+                    nn.Dropout2d(dropout),  # 2D dropout (matches paper specification)
                     nn.GroupNorm(
                         encoder_h // 2, encoder_h
                     ),  # Consider making num_groups configurable or ensure encoder_h is divisible by 2

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/models/eegminer.py

@@ -5,9 +5,10 @@
 * https://www.ipo.gov.uk/p-ipsum/Case/ApplicationNumber/GB2113420.0
 """
 
-# Authors: Sarthak Tayal <sarthaktayal2@gmail.com>
+from functools import partial
 
 import torch
+from einops.layers.torch import Rearrange
 from torch import nn
 
 import braindecode.functional as F
@@ -18,36 +19,6 @@ from braindecode.modules import GeneralizedGaussianFilter
 _eeg_miner_methods = ["mag", "corr", "plv"]
 
 
-class _EEGMinerFeatures(nn.Module):
-    def __init__(self, method, n_chans, n_filters, n_times):
-        super().__init__()
-        self.method = method
-        self.n_chans = n_chans
-        self.n_filters = n_filters
-        self.n_times = n_times
-        triu = torch.triu_indices(n_chans, n_chans, 1)
-        self.register_buffer("triu_row", triu[0], persistent=False)
-        self.register_buffer("triu_col", triu[1], persistent=False)
-
-    def forward(self, x):
-        if self.method == "mag":
-            return torch.sqrt((x * x).mean(dim=-1))
-        if self.method == "corr":
-            x = x.reshape(
-                x.shape[0], self.n_chans, self.n_filters, self.n_times
-            ).transpose(-3, -2)
-            x = (x - x.mean(dim=-1, keepdim=True)) / torch.sqrt(
-                x.var(dim=-1, keepdim=True) + 1e-6
-            )
-            x = torch.matmul(x, x.transpose(-2, -1)) / x.shape[-1]
-            x = x.permute(0, 2, 3, 1).abs()
-        else:
-            x = x.transpose(-4, -3)
-            x = F.plv_time(x, forward_fourier=False)
-            x = x.permute(0, 2, 3, 1)
-        return x[:, self.triu_row, self.triu_col, :]
-
-
 class EEGMiner(EEGModuleMixin, nn.Module):
     r"""EEGMiner from Ludwig et al (2024) [eegminer]_.
 
@@ -197,18 +168,24 @@ class EEGMiner(EEGModuleMixin, nn.Module):
         )
 
         # Forward method
-        self.ensure_dim = nn.Identity()
         if self.method == "mag":
+            self.method_forward = self._apply_mag_forward
             self.n_features = self.n_chans * self.n_filters
+            self.ensure_dim = nn.Identity()
         elif self.method == "corr":
+            self.method_forward = partial(
+                self._apply_corr_forward,
+                n_chans=self.n_chans,
+                n_filters=self.n_filters,
+                n_times=self.n_times,
+            )
             self.n_features = self.n_filters * self.n_chans * (self.n_chans - 1) // 2
+            self.ensure_dim = nn.Identity()
         elif self.method == "plv":
-            self.ensure_dim = nn.Unflatten(-1, (1, self.n_times))
+            self.method_forward = partial(self._apply_plv, n_chans=self.n_chans)
+            self.ensure_dim = Rearrange("... d -> ... 1 d")
             self.n_features = (self.n_filters * self.n_chans * (self.n_chans - 1)) // 2
 
-        self.feature_layer = _EEGMinerFeatures(
-            self.method, self.n_chans, self.n_filters, self.n_times
-        )
         self.flatten_layer = nn.Flatten()
         # Classifier
         self.batch_layer = nn.BatchNorm1d(self.n_features, affine=False)
@@ -218,17 +195,64 @@ class EEGMiner(EEGModuleMixin, nn.Module):
     @_disable_batch_norm_training_if_batch_size_one
     def forward(self, x):
         """x: (batch, electrodes, time)"""
+        batch = x.shape[0]
         x = self.ensure_dim(x)
         # Apply Gaussian filters in frequency domain
         # x -> (batch, electrodes * filters, time)
         x = self.filter(x)
 
-        x = self.feature_layer(x)
+        x = self.method_forward(x=x, batch=batch)
         # Classifier
         # Note that the order of dimensions before flattening the feature vector is important
         # for attributing feature weights during interpretation.
-        x = x.reshape(x.shape[0], self.n_features)
+        x = x.reshape(batch, self.n_features)
         x = self.batch_layer(x)
         x = self.final_layer(x)
 
         return x
+
+    @staticmethod
+    def _apply_mag_forward(x, batch=None):
+        # Signal magnitude
+        x = x * x
+        x = x.mean(dim=-1)
+        x = torch.sqrt(x)
+        return x
+
+    @staticmethod
+    def _apply_corr_forward(
+        x, batch, n_chans, n_filters, n_times, epilson: float = 1e-6
+    ):
+        x = x.reshape(batch, n_chans, n_filters, n_times).transpose(-3, -2)
+        x = (x - x.mean(dim=-1, keepdim=True)) / torch.sqrt(
+            x.var(dim=-1, keepdim=True) + epilson
+        )
+        x = torch.matmul(x, x.transpose(-2, -1)) / x.shape[-1]
+        # Original tensor shape: [batch, n_filters, chans, chans]
+        x = x.permute(0, 2, 3, 1)
+        # New tensor shape: [batch, chans, chans, n_filters]
+        # move filter channels to the end
+        x = x.abs()
+
+        # Get upper triu of symmetric connectivity matrix
+        triu = torch.triu_indices(n_chans, n_chans, 1)
+        x = x[:, triu[0], triu[1], :]
+
+        return x
+
+    @staticmethod
+    def _apply_plv(x, n_chans, batch=None):
+        # Compute PLV connectivity
+        # x -> (batch, electrodes, electrodes, filters)
+        x = x.transpose(-4, -3)  # swap electrodes and filters
+        # adjusting to compute the plv
+        x = F.plv_time(x, forward_fourier=False)
+        # batch, number of filters, connectivity matrix
+        # [batch, n_filters, chans, chans]
+        x = x.permute(0, 2, 3, 1)
+        # [batch, chans, chans, n_filters]
+
+        # Get upper triu of symmetric connectivity matrix
+        triu = torch.triu_indices(n_chans, n_chans, 1)
+        x = x[:, triu[0], triu[1], :]
+        return x

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/models/summary.csv

@@ -58,6 +58,5 @@ InterpolatedLaBraM,General,"Prediction, Embedding",200,"chs_info, n_outputs, n_t
 InterpolatedSignalJEPA,"Motor Imagery, ERP, SSVEP",Embedding,128,"chs_info",3456882,"InterpolatedSignalJEPA(chs_info=<user>)","Convolution,Channel,Foundation Model",EEG
 TCFormer,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",77820,"TCFormer(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer",EEG
 EEGDINO,General,Classification,200,"n_chans, n_outputs, n_times",4539698,"EEGDINO(n_chans=19, n_outputs=4, n_times=1000)","Foundation Model,Attention/Transformer",EEG
-DANCE,General,Event Detection,200,"n_chans, n_outputs, n_times, sfreq, chs_info",29813542,"DANCE(n_chans=19, n_outputs=4, n_times=6400, sfreq=200, input_window_seconds=32.0, chs_info=[{'ch_name': f'E{i+1}', 'kind': 'eeg', 'loc': numpy.random.default_rng(2024).random(12)} for i in range(19)])","Convolution,Attention/Transformer,Channel",EEG
 STEEGFormer,General,Classification,250,"n_chans, n_outputs, n_times",25305604,"STEEGFormer(n_chans=22, n_outputs=4, n_times=1000)","Attention/Transformer,Foundation Model",EEG
 MVPFormer,General,"Prediction, Embedding",512,"n_chans, n_outputs, n_times, sfreq",3356978,"MVPFormer(n_chans=22, n_outputs=4, n_times=2000, sfreq=100)","Foundation Model,Attention/Transformer",EEG

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/models/tcn.py

@@ -1,6 +1,5 @@
 # Authors: Patryk Chrabaszcz
 #          Lukas Gemein <l.gemein@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD-3
 import torch
@@ -237,9 +236,7 @@ class _TemporalBlock(nn.Module):
         )
         self.chomp1 = Chomp1d(padding)
         self.relu1 = activation()
-        # activations are (batch, channels, times), dropout2d only masks
-        # channel-wise there through a deprecated fallback
-        self.dropout1 = nn.Dropout1d(drop_prob)
+        self.dropout1 = nn.Dropout2d(drop_prob)
 
         self.conv2 = weight_norm(
             nn.Conv1d(
@@ -253,7 +250,7 @@ class _TemporalBlock(nn.Module):
         )
         self.chomp2 = Chomp1d(padding)
         self.relu2 = activation()
-        self.dropout2 = nn.Dropout1d(drop_prob)
+        self.dropout2 = nn.Dropout2d(drop_prob)
 
         self.downsample = (
             nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/models/util.py

@@ -484,25 +484,6 @@ models_mandatory_parameters: list[
     ),
     ("DGCNN", ["n_chans", "n_outputs", "n_times", "chs_info"], None),
     ("EEGDINO", ["n_chans", "n_outputs", "n_times"], None),
-    (
-        "DANCE",
-        ["n_outputs", "n_chans", "n_times", "sfreq", "chs_info"],
-        {
-            "n_chans": 19,
-            "n_times": 6400,  # 32 s @ 200 Hz, above the depth-10 stack minimum
-            "sfreq": 200.0,
-            "input_window_seconds": 32.0,
-            "n_outputs": 4,
-            "chs_info": [
-                {
-                    "ch_name": f"E{i + 1}",
-                    "kind": "eeg",
-                    "loc": _rng.random(12),
-                }
-                for i in range(19)
-            ],
-        },
-    ),
 ]
 
 ################################################################
@@ -517,9 +498,6 @@ non_classification_models = [
     # Emits a (batch, T_out, vocab) sequence for CTC, not class logits.
     "MetaNeuromotorHand",
     "EMG2QwertyNet",
-    # forward returns (batch, num_latents, n_outputs) dense per-token logits,
-    # not class logits.
-    "DANCE",
 ]
 
 ################################################################

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/modules/__init__.py

@@ -23,7 +23,6 @@ from .convolution import (
     Conv2dWithConstraint,
     DepthwiseConv2d,
 )
-from .dance_modules import Perceiver, SimpleConv
 from .filter import FilterBankLayer, GeneralizedGaussianFilter
 from .interpolation import ChannelInterpolationLayer
 from .layers import (
@@ -79,10 +78,6 @@ __all__ = [
     "Conv1dWithConstraint",
     "Conv2dWithConstraint",
     "DepthwiseConv2d",
-    "ChannelMerger",
-    "FourierEmb",
-    "Perceiver",
-    "SimpleConv",
     "FilterBankLayer",
     "GeneralizedGaussianFilter",
     "Chomp1d",
@@ -93,6 +88,8 @@ __all__ = [
     "TimeDistributed",
     "LinearWithConstraint",
     "MaxNormLinear",
+    "ChannelMerger",
+    "FourierEmb",
     "MaxNorm",
     "MaxNormParametrize",
     "LogPowerLayer",

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/modules/filter.py

@@ -1,7 +1,3 @@
-# Authors: Sarthak Tayal <sarthaktayal2@gmail.com>
-#
-# License: BSD-3
-
 from __future__ import annotations
 
 from typing import Optional
@@ -641,22 +637,20 @@ class GeneralizedGaussianFilter(nn.Module):
             The constructed filters with shape `(out_channels, freq_bins, 2)`.
 
         """
-        min_f_mean = self.clamp_f_mean[0] / (self.sample_rate / 2)
-        max_f_mean = self.clamp_f_mean[1] / (self.sample_rate / 2)
-        f_mean = torch.clamp(self.f_mean, min=min_f_mean, max=max_f_mean)
-        bandwidth = torch.clamp(
-            self.bandwidth, min=1.0 / (self.sample_rate / 2), max=1.0
+        # Clamp parameters
+        self.f_mean.data = torch.clamp(
+            self.f_mean.data,
+            min=self.clamp_f_mean[0] / (self.sample_rate / 2),
+            max=self.clamp_f_mean[1] / (self.sample_rate / 2),
         )
-        shape = torch.clamp(self.shape, min=2.0, max=3.0)
-
-        if not torch.jit.is_scripting():
-            self.f_mean.data = f_mean
-            self.bandwidth.data = bandwidth
-            self.shape.data = shape
+        self.bandwidth.data = torch.clamp(
+            self.bandwidth.data, min=1.0 / (self.sample_rate / 2), max=1.0
+        )
+        self.shape.data = torch.clamp(self.shape.data, min=2.0, max=3.0)
 
         # Create magnitude response with gain=1 -> (channels, freqs)
         mag_response = self.exponential_power(
-            self.n_range, f_mean, bandwidth, shape * 8 - 14
+            self.n_range, self.f_mean, self.bandwidth, self.shape * 8 - 14
         )
         mag_response = mag_response / mag_response.max(dim=-1, keepdim=True)[0]
 

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/regressor.py

@@ -3,7 +3,6 @@
 #          Lukas Gemein <l.gemein@gmail.com>
 #          Bruno Aristimunha <b.aristimunha@gmail.com>
 #          Pierre Guetschel <pierre.guetschel@gmail.com>
-#          Sarthak Tayal <sarthaktayal2@gmail.com>
 #
 # License: BSD (3-clause)
 
@@ -140,7 +139,9 @@ class EEGRegressor(_EEGNeuralNet, NeuralNetRegressor):
             return y_pred
 
     def predict_trials(self, X, return_targets=True):
-        """Create trialwise predictions from a cropped dataset.
+        """Create trialwise predictions and optionally also return trialwise.
+
+        labels from cropped dataset.
 
         Parameters
         ----------
@@ -155,11 +156,10 @@ class EEGRegressor(_EEGNeuralNet, NeuralNetRegressor):
             3-dimensional array (n_trials x n_classes x n_predictions), where
             the number of predictions depend on the chosen window size and the
             receptive field of the network.
-        trial_targets : np.ndarray
-            Ground-truth targets from the dataset in a 2-dimensional array
-            (n_trials x n_targets). Only returned when ``return_targets=True``.
-            The number of targets depends on the decoding paradigm and can be
-            either a single value, multiple values, or a sequence.
+        trial_labels : np.ndarray
+            2-dimensional array (n_trials x n_targets) where the number of
+            targets depends on the decoding paradigm and can be either a single
+            value, multiple values, or a sequence.
         """
         if not self.cropped:
             warnings.warn(

{braindecode-1.7.0 → braindecode-1.7.0.dev174519142}/braindecode/training/__init__.py

@@ -2,19 +2,17 @@
 Functionality for skorch-based training.
 """
 
-from .losses import CroppedLoss, DanceLoss, TimeSeriesLoss, mixup_criterion
+from .losses import CroppedLoss, TimeSeriesLoss, mixup_criterion
 from .scoring import (
     CroppedTimeSeriesEpochScoring,
     CroppedTrialEpochScoring,
     PostEpochTrainScoring,
-    f1_event,
     predict_trials,
     trial_preds_from_window_preds,
 )
 
 __all__ = [
     "CroppedLoss",
-    "DanceLoss",
     "mixup_criterion",
     "TimeSeriesLoss",
     "CroppedTrialEpochScoring",
@@ -22,5 +20,4 @@ __all__ = [
     "CroppedTimeSeriesEpochScoring",
     "trial_preds_from_window_preds",
     "predict_trials",
-    "f1_event",
 ]

braindecode-1.7.0.dev174519142/braindecode/training/losses.py

@@ -0,0 +1,105 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#          Maciej Sliwowski <maciek.sliwowski@gmail.com>
+#          Mohammed Fattouh <mo.fattouh@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+from torch import nn
+
+
+class CroppedLoss(nn.Module):
+    """Compute Loss after averaging predictions across time.
+    Assumes predictions are in shape:
+    n_batch size x n_classes x n_predictions (in time)"""
+
+    def __init__(self, loss_function):
+        super().__init__()
+        self.loss_function = loss_function
+
+    def forward(self, preds, targets):
+        """Forward pass.
+
+        Parameters
+        ----------
+        preds: torch.Tensor
+            Model's prediction with shape (batch_size, n_classes, n_times).
+        targets: torch.Tensor
+            Target labels with shape (batch_size, n_classes, n_times).
+        """
+        avg_preds = torch.mean(preds, dim=2)
+        avg_preds = avg_preds.squeeze(dim=1)
+        return self.loss_function(avg_preds, targets)
+
+
+class TimeSeriesLoss(nn.Module):
+    """Compute Loss between timeseries targets and predictions.
+    Assumes predictions are in shape:
+    n_batch size x n_classes x n_predictions (in time)
+    Assumes targets are in shape:
+    n_batch size x n_classes x window_len (in time)
+    If the targets contain NaNs, the NaNs will be masked out and the loss will be only computed for
+    predictions valid corresponding to valid target values."""
+
+    def __init__(self, loss_function):
+        super().__init__()
+        self.loss_function = loss_function
+
+    def forward(self, preds, targets):
+        """Forward pass.
+
+        Parameters
+        ----------
+        preds: torch.Tensor
+            Model's prediction with shape (batch_size, n_classes, n_times).
+        targets: torch.Tensor
+            Target labels with shape (batch_size, n_classes, n_times).
+        """
+        n_preds = preds.shape[-1]
+        # slice the targets to fit preds shape
+        targets = targets[:, :, -n_preds:]
+        # create valid targets mask
+        mask = ~torch.isnan(targets)
+        # select valid targets that have a matching predictions
+        masked_targets = targets[mask]
+        masked_preds = preds[mask]
+        return self.loss_function(masked_preds, masked_targets)
+
+
+def mixup_criterion(preds, target):
+    """Implements loss for Mixup for EEG data. See [1]_.
+
+    Implementation based on [2]_.
+
+    Parameters
+    ----------
+    preds : torch.Tensor
+        Predictions from the model.
+    target : torch.Tensor | list of torch.Tensor
+        For predictions without mixup, the targets as a tensor. If mixup has
+        been applied, a list containing the targets of the two mixed
+        samples and the mixing coefficients as tensors.
+
+    Returns
+    -------
+    loss : float
+        The loss value.
+
+    References
+    ----------
+    .. [1] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz
+       mixup: Beyond Empirical Risk Minimization
+       Online: https://arxiv.org/abs/1710.09412
+    .. [2] https://github.com/facebookresearch/mixup-cifar10/blob/master/train.py
+    """
+    if len(target) == 3:
+        # unpack target
+        y_a, y_b, lam = target
+        # compute loss per sample
+        loss_a = torch.nn.functional.nll_loss(preds, y_a, reduction="none")
+        loss_b = torch.nn.functional.nll_loss(preds, y_b, reduction="none")
+        # compute weighted mean
+        ret = torch.mul(lam, loss_a) + torch.mul(1 - lam, loss_b)
+        return ret.mean()
+    else:
+        return torch.nn.functional.nll_loss(preds, target)