fix #391 (comment)

spikingjelly/activation_based/layer.py

@@ -1872,6 +1872,7 @@ def __init__(self, alpha: float, v_th: float, *args, **kwargs):
         torch.nn.init.constant_(self.weight, alpha * v_th)
 
     def forward(self, x_seq):
+        assert self.step_mode == 'm', "ThresholdDependentBatchNormBase can only be used in the multi-step mode!"
         return functional.seq_to_ann_forward(x_seq, super().forward)
 
 
@@ -1907,6 +1908,10 @@ def __init__(self, alpha: float, v_th: float, *args, **kwargs):
         """
         super().__init__(alpha, v_th, *args, **kwargs)
 
+    def _check_input_dim(self, input):
+        assert input.dim() == 4 - 1  # [T * N, C, L]
+
+
 
 class ThresholdDependentBatchNorm2d(_ThresholdDependentBatchNormBase):
     def __init__(self, alpha: float, v_th: float, *args, **kwargs):
@@ -1940,6 +1945,8 @@ def __init__(self, alpha: float, v_th: float, *args, **kwargs):
         """
         super().__init__(alpha, v_th, *args, **kwargs)
 
+    def _check_input_dim(self, input):
+        assert input.dim() == 5 - 1  # [T * N, C, H, W]
 
 class ThresholdDependentBatchNorm3d(_ThresholdDependentBatchNormBase):
     def __init__(self, alpha: float, v_th: float, *args, **kwargs):
@@ -1973,6 +1980,9 @@ def __init__(self, alpha: float, v_th: float, *args, **kwargs):
         """
         super().__init__(alpha, v_th, *args, **kwargs)
 
+    def _check_input_dim(self, input):
+        assert input.dim() == 6 - 1  # [T * N, C, H, W, D]
+
 
 class TemporalWiseAttention(nn.Module, base.MultiStepModule):
     def __init__(self, T: int, reduction: int = 16, dimension: int = 4):