braindecode 1.8.1.dev183280391__tar.gz → 1.8.1.dev183301950__tar.gz

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/NOTICE.txt

@@ -54,6 +54,7 @@ The following components are licensed under the MIT License:
 - `docs/_static/model/seizure_transformer_arch.png`
 - `braindecode/models/guetschel2026.py`
 - `braindecode/models/tmsanet.py`
+- `braindecode/models/epint.py`
 
 The `brainomni.py` file is a derivative of `OpenTSLab/BrainOmni` (MIT), whose
 SEANet codec derives from Meta's `facebookresearch/encodec` (MIT). The

{braindecode-1.8.1.dev183280391/braindecode.egg-info → braindecode-1.8.1.dev183301950}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev183280391
+Version: 1.8.1.dev183301950
 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.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/datasets/base.py

@@ -1518,7 +1518,9 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
         metadata : pd.DataFrame
             DataFrame containing as many rows as there are windows in the
             BaseConcatDataset, with the metadata and description information
-            for each window.
+            for each window, and an ``i_dataset`` column giving the position
+            in ``datasets`` of the dataset (recording) each window comes from.
+            Samplers use it to keep recordings apart.
         """
         if not all(
             [
@@ -1544,9 +1546,15 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
                     "Dataset descriptions cannot contain the reserved window "
                     "metadata key 'i_trial_in_dataset'."
                 )
+            if "i_dataset" in df.columns or "i_dataset" in ds.description:
+                raise ValueError(
+                    "Window metadata and dataset descriptions cannot contain the "
+                    "key 'i_dataset', which get_metadata() reserves for the "
+                    "position of the dataset each window comes from. Rename it."
+                )
 
         all_dfs = list()
-        for ds in self.datasets:
+        for i_ds, ds in enumerate(self.datasets):
             if hasattr(ds, "_windows") and ds._windows is not None:
                 df = ds._windows.metadata
             else:
@@ -1556,6 +1564,7 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
                 df = df.copy()
             for k, v in ds.description.items():
                 df[k] = v
+            df["i_dataset"] = i_ds
             all_dfs.append(df)
 
         return pd.concat(all_dfs)

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/functional/functions.py

@@ -91,16 +91,27 @@ def fft_conv1d(
     n_fft = 2 * _next_fast_len((n_times + kernel_size) // 2)
     xs = spectral_input(x)
     # conv1d is a cross-correlation: a convolution with the flipped kernel
-    x_f = torch.fft.rfft(xs, n=n_fft).permute(2, 0, 1)  # (freq, batch, in)
+    x_f = torch.fft.rfft(xs, n=n_fft)  # (batch, in, freq)
     w_f = torch.fft.rfft(weight.to(xs).flip(-1), n=n_fft).permute(2, 1, 0)
     # Per frequency (batch, in) @ (in, out) in complex numbers, as one real
-    # matmul (a complex one is slower on CPU): [Re x; Im x] @ [Re w, Im w]
-    # (columns interleaved) holds the four real products.
-    p = torch.cat([x_f.real, x_f.imag], 1) @ torch.view_as_real(w_f).flatten(2)
-    p = p.unflatten(1, [2, x.shape[0]]).unflatten(3, [-1, 2])  # x re/im, w re/im
-    y_f = torch.complex(
-        p[:, 0, :, :, 0] - p[:, 1, :, :, 1], p[:, 0, :, :, 1] + p[:, 1, :, :, 0]
-    )
+    # matmul (a complex one is slower on CPU).
+    needs_grad = torch.is_grad_enabled() and (x.requires_grad or weight.requires_grad)
+    if needs_grad or x.shape[0] >= 2 * x.shape[1]:
+        # [Re x, Im x] @ [[Re w, Im w], [-Im w, Re w]] (rows of i * w): the
+        # output columns are the (Re, Im) pairs of the result, read as complex
+        # without a copy (and without a costly backward).
+        w_mat = torch.view_as_real(torch.stack([w_f, w_f * 1j], 1)).flatten(1, 2)
+        y_f = torch.view_as_real(x_f).permute(2, 0, 3, 1).flatten(2) @ w_mat.flatten(2)
+        y_f = torch.view_as_complex(y_f.unflatten(-1, [-1, 2]))
+    else:
+        # Forward only with fewer rows than twice the input channels: the
+        # doubled weight costs more than combining [Re x; Im x] @ [Re w, Im w].
+        x_f = x_f.permute(2, 0, 1)
+        p = torch.cat([x_f.real, x_f.imag], 1) @ torch.view_as_real(w_f).flatten(2)
+        p = p.unflatten(1, [2, x.shape[0]]).unflatten(3, [-1, 2])
+        y_f = torch.complex(
+            p[:, 0, :, :, 0] - p[:, 1, :, :, 1], p[:, 0, :, :, 1] + p[:, 1, :, :, 0]
+        )
     y = torch.fft.irfft(y_f.permute(1, 2, 0), n=n_fft)
     # "same" output t is sample t + kernel_size // 2 of the full convolution
     y = y[..., kernel_size // 2 : kernel_size // 2 + n_times]

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/__init__.py

@@ -36,6 +36,7 @@ from .eegsimpleconv import EEGSimpleConv
 from .eegsym import EEGSym
 from .eegtcnet import EEGTCNet
 from .emg2qwerty import EMG2QwertyNet
+from .epint import EpiNT
 from .fbcnet import FBCNet
 from .fblightconvnet import FBLightConvNet
 from .fbmsnet import FBMSNet
@@ -134,6 +135,7 @@ __all__ = [
     "EEGTCNet",
     "DGCNN",
     "EMG2QwertyNet",
+    "EpiNT",
     "NeuroPose",
     "SensingDynamics",
     "VEMG2Pose",

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/axon.py

@@ -14,6 +14,7 @@ from torch import nn
 from braindecode.functional import sinusoidal_positional_encoding
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import warn_if_sfreq_differs
+from braindecode.modules.attention import _scaled_dot_product_attention
 from braindecode.modules.channels import _resolve
 
 
@@ -462,11 +463,33 @@ class _AxisAttention(nn.Module):
             .unbind(0)
         )
         drop = self.att_drop_prob if self.training else 0.0
-        attended = F.scaled_dot_product_attention(
+        attended = _scaled_dot_product_attention(
             query, key, value, attn_mask=mask, dropout_p=drop
         )
         return attended.transpose(1, 2).reshape(n_seq, length, -1)
 
+    def _attend_two_scales(
+        self, projected: torch.Tensor, band: torch.Tensor, weights: torch.Tensor
+    ) -> torch.Tensor:
+        """Banded and full attention from one score matrix, mixed per query.
+
+        ``weights`` (N, L, 2) are the scale-gate weights. Both softmaxes share
+        the scores and one product with the values serves the mixture.
+        """
+        n_seq, length, _ = projected.shape
+        query, key, value = (
+            projected.reshape(n_seq, length, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+            .unbind(0)
+        )
+        scores = torch.matmul(query, key.transpose(-2, -1)) * self.head_dim**-0.5
+        weights = weights.reshape(n_seq, 1, length, 2)
+        probs = weights[..., 0:1] * scores.masked_fill(~band, float("-inf")).softmax(
+            -1
+        ) + weights[..., 1:2] * scores.softmax(-1)
+        attended = torch.matmul(probs, value)
+        return attended.transpose(1, 2).reshape(n_seq, length, -1)
+
     def forward(self, x: torch.Tensor, band: Optional[torch.Tensor]) -> torch.Tensor:
         batch, n_chans, n_patches, dim = x.shape
         axis_weights = self.axis_gate(x)  # (B, C, P, 2)
@@ -476,15 +499,26 @@ class _AxisAttention(nn.Module):
         # the restricted branch is the full one; attend once unless the two
         # branches draw their own attention dropout.
         per_channel = self.temporal_qkv(x.reshape(batch * n_chans, n_patches, dim))
-        restricted = self.temporal_out(self._attend(per_channel, band)).reshape(
-            batch, n_chans, n_patches, dim
-        )
-        full = restricted
-        if band is not None or (self.training and self.att_drop_prob > 0):
-            full = self.temporal_out(self._attend(per_channel, None)).reshape(
+        dropout = self.training and self.att_drop_prob > 0
+        if band is not None and not dropout:
+            # The scale-gate weights sum to one, so the two branches share the
+            # scores, the product with the values and temporal_out.
+            attended = self._attend_two_scales(per_channel, band, scale_weights)
+            temporal = self.temporal_out(attended).reshape(
                 batch, n_chans, n_patches, dim
             )
-        temporal = scale_weights[..., 0:1] * restricted + scale_weights[..., 1:2] * full
+        else:
+            restricted = self.temporal_out(self._attend(per_channel, band)).reshape(
+                batch, n_chans, n_patches, dim
+            )
+            full = restricted
+            if dropout:
+                full = self.temporal_out(self._attend(per_channel, None)).reshape(
+                    batch, n_chans, n_patches, dim
+                )
+            temporal = (
+                scale_weights[..., 0:1] * restricted + scale_weights[..., 1:2] * full
+            )
 
         # Spatial path: one sequence of C tokens per time step.
         per_step = self.spatial_qkv(

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/brainmodule.py

@@ -20,6 +20,7 @@ from torch.nn import functional as F
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import has_valid_locations, positions_from_chs_info
 from braindecode.modules import ChannelMerger
+from braindecode.modules.attention import _hpu_lazy
 from braindecode.modules.layers import SubjectLayers
 
 __all__ = ["BrainModule"]
@@ -712,7 +713,7 @@ class _ConvSequence(nn.Module):
                         nn.Conv1d(
                             chout, 2 * chout, 1 + 2 * glu_context, padding=glu_context
                         ),
-                        nn.GLU(dim=1),
+                        _GLU(dim=1),
                     )
                 )
             else:
@@ -727,6 +728,17 @@ class _ConvSequence(nn.Module):
         return x
 
 
+class _GLU(nn.GLU):
+    """``nn.GLU`` that Intel Gaudi's lazy mode can compile."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not torch.jit.is_scripting() and x.device.type == "hpu" and _hpu_lazy():
+            # Its graph compiler fails on ``aten.glu`` after a convolution.
+            a, b = x.chunk(2, dim=self.dim)
+            return a * torch.sigmoid(b)
+        return F.glu(x, self.dim)
+
+
 def _pad_multiple(x: torch.Tensor, base: int) -> torch.Tensor:
     """Pad tensor to be a multiple of base."""
     length = x.shape[-1]

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/codebrain.py

@@ -17,6 +17,7 @@ from torch.nn.utils.parametrizations import weight_norm
 from braindecode.functional import spectral_input
 from braindecode.functional.functions import _next_fast_len
 from braindecode.models.base import EEGModuleMixin
+from braindecode.modules.attention import _hpu_lazy
 
 
 class CodeBrain(EEGModuleMixin, nn.Module):
@@ -831,10 +832,20 @@ class _ResidualGroup(nn.Module):
     def forward(self, input_data: torch.Tensor) -> torch.Tensor:
         noise = input_data
         h = noise
-        skip = torch.zeros((), dtype=h.dtype, device=h.device)
-        for block in self.residual_blocks:
-            h, skip_n = block((h, noise))
-            skip = skip_n + skip
+        if not torch.jit.is_scripting() and h.device.type == "hpu" and _hpu_lazy():
+            # Gaudi's lazy-mode graph compiler fails on the backward of the
+            # sequential skip sum (one skip_conv dedw node lands in two
+            # pipeline bundles); one n-ary stack-sum reduction compiles.
+            skip_list = []
+            for block in self.residual_blocks:
+                h, skip_n = block((h, noise))
+                skip_list.append(skip_n)
+            skip = torch.stack(skip_list, dim=0).sum(dim=0)
+        else:
+            skip = torch.zeros((), dtype=h.dtype, device=h.device)
+            for block in self.residual_blocks:
+                h, skip_n = block((h, noise))
+                skip = skip_n + skip
 
         return skip * math.sqrt(1.0 / self.num_res_layers)
 

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/ctnet.py

@@ -378,7 +378,15 @@ class _PatchEmbeddingEEGNet(nn.Module):
         Tensor
             Embedded patches of shape (batch_size, num_patches, embedding_dim).
         """
-        x = self.eegnet_module(x)
+        layers = self.eegnet_module
+        x = layers[1](layers[0](x))
+        # The depthwise (n_chans, 1) convolution as a matmul over the channel
+        # axis: same values, a faster backward than the grouped convolution's.
+        b, f, c, t = x.shape
+        x = torch.matmul(layers[2].weight.view(f, -1, c), x).view(b, -1, 1, t)
+        for i, layer in enumerate(layers):
+            if i > 2:
+                x = layer(x)
         x = self.projection(x)
         return x
 

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/eegconformer.py

@@ -4,6 +4,7 @@
 import warnings
 
 import torch
+import torch.nn.functional as F
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
@@ -346,7 +347,19 @@ class _PatchEmbedding(nn.Module):
         )
 
     def forward(self, x: Tensor) -> Tensor:
-        x = self.shallownet(x)
+        # Nothing runs between the temporal and the spatial convolution, so
+        # they run as one convolution with the composed weight and bias (as
+        # CombinedConv does), without the (batch, filters, chans, time)
+        # intermediate.
+        conv_time, conv_spat = self.shallownet[0], self.shallownet[1]
+        bias_time, bias_spat = conv_time.bias, conv_spat.bias
+        assert bias_time is not None and bias_spat is not None
+        weight = (conv_time.weight * conv_spat.weight.permute(1, 0, 2, 3)).sum(0)
+        bias = conv_spat.weight.sum((2, 3)) @ bias_time + bias_spat
+        x = F.conv1d(x.squeeze(1), weight, bias).unsqueeze(2)
+        for i, layer in enumerate(self.shallownet):
+            if i > 1:
+                x = layer(x)
         x = self.projection(x)
         return x
 

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/eegnex.py

@@ -358,10 +358,14 @@ class EEGNeX(EEGModuleMixin, nn.Module):
 class _TemporalConv(nn.Conv2d):
     """``(1, k)`` Conv2d with "same" output length; on CPU and HPU, through
     :func:`~braindecode.functional.fft_conv1d` where
-    :func:`~braindecode.functional.prefer_fft_conv` expects it to be faster."""
+    :func:`~braindecode.functional.prefer_fft_conv` expects it to be faster.
+    On HPU only when gradients are computed: the FFT pays off in the
+    backward, and Gaudi's direct forward alone is cheaper."""
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if not prefer_fft_conv(x, self.kernel_size[1]):
+        if not prefer_fft_conv(x, self.kernel_size[1]) or (
+            x.device.type == "hpu" and not torch.is_grad_enabled()
+        ):
             return self._conv_forward(x, self.weight, self.bias)
         b, _, h, t = x.shape  # the conv1d of each of the b * h rows
         x = x.transpose(1, 2).reshape(b * h, -1, t)

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/eegpt.py

@@ -13,6 +13,7 @@ from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import DropPath
+from braindecode.modules.attention import _scaled_dot_product_attention
 from braindecode.modules.convolution import Conv1dWithConstraint
 from braindecode.modules.linear import LinearWithConstraint
 from braindecode.util import resolve_montage_name
@@ -949,14 +950,17 @@ class _Attention(nn.Module):
         # 3. Flash Attention
         # Use PyTorch's optimized scaled_dot_product_attention (SDPA)
         # which internally uses FlashAttention or EfficientAttention kernels.
-        y = torch.nn.functional.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=None,
-            dropout_p=self.attn_drop.p if self.training else 0.0,
-            is_causal=self.is_causal,
-        )
+        dropout_p = self.attn_drop.p if self.training else 0.0
+        if self.is_causal:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=True
+            )
+        else:
+            # Gaudi's fused kernel also in lazy mode: faster for EEGPT's
+            # 23-token sequences except on the shortest inputs (256 samples).
+            y = _scaled_dot_product_attention(
+                q, k, v, dropout_p=dropout_p, fused_in_lazy=True
+            )
 
         # Reshape back: (batch, num_heads, seq_len, head_dim) -> (batch, seq_len, embed_dim)
         x = y.transpose(1, 2).flatten(2)

{braindecode-1.8.1.dev183280391 → braindecode-1.8.1.dev183301950}/braindecode/models/eegsym.py

@@ -457,6 +457,58 @@ class _SpatialConv3d(nn.Conv3d):
         return y.reshape(b, -1, z, t, y.shape[-1])
 
 
+class _BatchNorm3d(nn.BatchNorm3d):
+    """``nn.BatchNorm3d`` with running statistics (as built here); on HPU it
+    normalises the ``(batch, channels, -1)`` view, same statistics and
+    buffers: Gaudi's 5-D batch norm backward is 2-7x slower than its 3-D one
+    on these activations. Elsewhere the same ``batch_norm`` call as torch's."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training and self.num_batches_tracked is not None:
+            self.num_batches_tracked.add_(1)
+        y = nn.functional.batch_norm(
+            x.flatten(2) if x.device.type == "hpu" else x,
+            self.running_mean,
+            self.running_var,
+            self.weight,
+            self.bias,
+            self.training,
+            self.momentum,  # type: ignore[arg-type]
+            self.eps,
+        )
+        return y.view_as(x)
+
+
+class _TemporalConv3d(nn.Conv3d):
+    """``nn.Conv3d`` with a ``(1, k, 1)`` kernel (padding on time only); on HPU
+    it runs as ``conv1d`` on a ``(batch * z * space, channels, time)`` copy of
+    the input, same parameters. Gaudi's conv3d along time, which is not the
+    innermost axis here, is 10-50x slower than that conv1d."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.time_padding = (
+            self.padding if isinstance(self.padding, str) else self.padding[1]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.device.type == "hpu":
+            return self._time_last(x)
+        return self._conv_forward(x, self.weight, self.bias)
+
+    def _time_last(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, z, t, s = x.shape
+        y = nn.functional.conv1d(
+            x.permute(0, 2, 4, 1, 3).reshape(b * z * s, c, t),
+            self.weight.flatten(2),
+            self.bias,
+            padding=self.time_padding,
+            groups=self.groups,
+        )
+        y = y.reshape(b, z, s, y.shape[1], y.shape[2])
+        return y.permute(0, 3, 1, 4, 2).contiguous()
+
+
 class _InceptionBlock(nn.Module):
     r"""Inception module used in EEGSym architecture.
 
@@ -498,13 +550,13 @@ class _InceptionBlock(nn.Module):
         for scale in scales_samples:
             self.temporal_convs.append(
                 nn.Sequential(
-                    nn.Conv3d(
+                    _TemporalConv3d(
                         in_channels=in_channels,
                         out_channels=filters_per_branch,
                         kernel_size=(1, scale, 1),
                         padding="same",
                     ),
-                    nn.BatchNorm3d(filters_per_branch, eps=1e-3, momentum=0.01),
+                    _BatchNorm3d(filters_per_branch, eps=1e-3, momentum=0.01),
                     activation,
                     nn.Dropout(drop_prob),
                 )
@@ -527,7 +579,7 @@ class _InceptionBlock(nn.Module):
                             groups=filters_per_branch * len(scales_samples),
                             bias=False,
                         ),
-                        nn.BatchNorm3d(
+                        _BatchNorm3d(
                             filters_per_branch * len(scales_samples),
                             eps=1e-3,
                             momentum=0.01,
@@ -598,13 +650,13 @@ class _ResidualBlock(nn.Module):
 
         # Temporal convolution
         self.temporal_conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=filters,
                 kernel_size=(1, kernel_size, 1),
                 padding="same",
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -618,7 +670,7 @@ class _ResidualBlock(nn.Module):
                 kernel_size=(1, 1, 1),
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -642,7 +694,7 @@ class _ResidualBlock(nn.Module):
                             groups=1,
                             bias=False,
                         ),
-                        nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+                        _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
                         activation,
                         nn.Dropout(drop_prob),
                     )
@@ -694,14 +746,14 @@ class _TemporalBlock(nn.Module):
         super().__init__()
 
         self.conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=filters,
                 kernel_size=(1, kernel_size, 1),
                 padding="same",
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -762,7 +814,7 @@ class _ChannelMergingBlock(nn.Module):
                         padding=(0, 0, 0),  # Valid padding
                         bias=False,
                     ),
-                    nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+                    _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
                     activation,
                     nn.Dropout(drop_prob),
                 )
@@ -779,7 +831,7 @@ class _ChannelMergingBlock(nn.Module):
                 padding=(0, 0, 0),
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -837,21 +889,21 @@ class _TemporalMergingBlock(nn.Module):
 
         # Residual convolution (collapses time dimension)
         self.residual_conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=in_channels,  # Same channels for residual
                 kernel_size=(1, self.temporal_kernel, 1),  # (Z, Time, Space)
                 padding=(0, 0, 0),  # Valid padding - reduces time to 1
                 bias=False,
             ),
-            nn.BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
 
         # Grouped convolution (collapses time dimension, doubles filters)
         self.grouped_conv = nn.Sequential(
-            nn.Conv3d(
+            _TemporalConv3d(
                 in_channels=in_channels,
                 out_channels=filters,  # Double the channels
                 kernel_size=(1, self.temporal_kernel, 1),  # (Z, Time, Space)
@@ -859,7 +911,7 @@ class _TemporalMergingBlock(nn.Module):
                 padding=(0, 0, 0),
                 bias=False,
             ),
-            nn.BatchNorm3d(filters, eps=1e-3, momentum=0.01),
+            _BatchNorm3d(filters, eps=1e-3, momentum=0.01),
             activation,
             nn.Dropout(drop_prob),
         )
@@ -908,7 +960,7 @@ class _OutputBlock(nn.Module):
                         padding=(0, 0, 0),
                         bias=False,
                     ),
-                    nn.BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
+                    _BatchNorm3d(in_channels, eps=1e-3, momentum=0.01),
                     activation,
                     nn.Dropout(drop_prob),
                 )