Repository: paul-krug/pytorch-tcn (language: python)

Issue:
Cannot use new/different activation function
I want to use [SiLu](https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html) as the activation function, but I keep getting errors from the `check_activation_arg` function when I pass the class
```
ValueError:
            The argument 'activation' must either be a valid string or
            a torch.nn.Module object, but <class '__main__.SiLU'> was passed,
            which is of type <class 'type'>.
```

and when I pass it as an object it throws this error  

```
    tcn_model = TCN(**model_hyperparams)
 pytorch_tcn/tcn.py", line 530, in __init__
    TemporalBlock(
 /pytorch_tcn/tcn.py", line 257, in __init__
    self.activation1 = self.activation()
                       ^^^^^^^^^^^^^^^^^
 /torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 /torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: SiLU.forward() missing 1 required positional argument: 'input'
```

I'm using "pytorch-tcn>=1.2.1", and python 3.12.6.

Right now I'm passing an instance of this wrapper-class as a workaround:

```python
class SiLU(torch.nn.Module):
    def __call__(self, *args, **kwargs):
        return torch.nn.SiLU()
```




Discussion / hints:

Fix patch:
diff --git a/pytorch_tcn/__init__.py b/pytorch_tcn/__init__.py
index 9f5082d..702d04c 100644
--- a/pytorch_tcn/__init__.py
+++ b/pytorch_tcn/__init__.py
@@ -2,4 +2,4 @@
 from pytorch_tcn.conv import TemporalConv1d
 from pytorch_tcn.conv import TemporalConvTranspose1d
 
-__version__ = '1.2.1'
\ No newline at end of file
+__version__ = '1.2.2.dev0'
\ No newline at end of file
diff --git a/pytorch_tcn/buffer.py b/pytorch_tcn/buffer.py
new file mode 100644
index 0000000..bb26126
--- /dev/null
+++ b/pytorch_tcn/buffer.py
@@ -0,0 +1,55 @@
+
+import torch
+
+from typing import Optional
+from typing import Union
+from typing import List
+from collections.abc import Iterable
+
+
+class BufferIO():
+    def __init__(
+            self,
+            in_buffers: Optional[ Iterable ] = None,
+            ):
+        if in_buffers is not None:
+            self.in_buffers_length = len( in_buffers )
+            self.in_buffers = iter( in_buffers )
+        else:
+            self.in_buffers_length = None
+            self.in_buffers = None
+        
+        self.out_buffers = []
+        return
+    
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        if self.in_buffers is not None:
+            return next( self.in_buffers )
+        else:
+            return None
+        
+    def append_out_buffer(
+            self,
+            x: torch.Tensor,
+            ):
+        self.out_buffers.append(x)
+        return
+        
+    def next_in_buffer(
+            self,
+            ):
+        return self.__next__()
+        
+    def step(self):
+        if len( self.out_buffers ) != self.in_buffers_length:
+            raise ValueError(
+                """
+                Number of out buffers does not match number of in buffers.
+                """
+                )
+        self.in_buffers = iter( self.out_buffers )
+        self.out_buffers = []
+        return
\ No newline at end of file
diff --git a/pytorch_tcn/conv.py b/pytorch_tcn/conv.py
index e78b48f..37840a5 100644
--- a/pytorch_tcn/conv.py
+++ b/pytorch_tcn/conv.py
@@ -4,13 +4,14 @@
 import torch.nn as nn
 import math
 
-# Padding modes
-PADDING_MODES = [
-    'zeros',
-    'reflect',
-    'replicate',
-    'circular',
-]
+from .pad import TemporalPad1d
+from .buffer import BufferIO
+
+from typing import Optional
+from typing import Union
+from typing import List
+    
+
 
 class TemporalConv1d(nn.Conv1d):
     def __init__(
@@ -76,115 +77,63 @@ def __init__(
                 """
                 )
 
-
         self.pad_len = (kernel_size - 1) * dilation
         self.causal = causal
         
-        if causal:
-            # Padding is only on the left side
-            self.left_pad = self.pad_len
-            self.right_pad = 0
-        else:
-            # Padding is on both sides
-            self.left_pad = self.pad_len // 2
-            self.right_pad = self.pad_len - self.left_pad
-        
-        if padding_mode == 'zeros':
-            self.padder = nn.ConstantPad1d(
-                ( self.left_pad, self.right_pad ),
-                0.0,
-                )
-        elif padding_mode == 'reflect':
-            self.padder = nn.ReflectionPad1d(
-                ( self.left_pad, self.right_pad ),
-                )
-        elif padding_mode == 'replicate':
-            self.padder = nn.ReplicationPad1d(
-                ( self.left_pad, self.right_pad ),
-                )
-        elif padding_mode == 'circular':
-            self.padder = nn.CircularPad1d(
-                ( self.left_pad, self.right_pad ),
-                )
-        else:
-            raise ValueError(
-                f"""
-                padding_mode must be one of {PADDING_MODES},
-                but got {padding_mode}.
-                """
-                )
-
-        # Buffer is used for streaming inference
-        if buffer is None:
-            buffer = torch.zeros(
-                1,
-                in_channels,
-                self.pad_len,
-                )
-        
-        # Register buffer as a persistent buffer which is available as self.buffer
-        self.register_buffer(
-            'buffer',
-            buffer,
+        self.padder = TemporalPad1d(
+            padding = self.pad_len,
+            in_channels = in_channels,
+            buffer = buffer,
+            padding_mode = padding_mode,
+            causal = causal,
             )
         
         return
     
-    def _forward(self, x):
-        x = self.padder(x)
-        x = super().forward(x)
-        return x
+    # In pytorch-tcn >= 1.2.2, buffer is moved to TemporalPad1d
+    # We keep the property for backwards compatibility, e.g. in
+    # case one wants to load old model weights.
+    @property
+    def buffer(self):
+        return self.padder.buffer
+    
+    @buffer.setter
+    def buffer(self, value):
+        self.padder.buffer = value
+        return
 
     def forward(
             self,
-            x,
-            inference=False,
-            in_buffer=None,
+            x: torch.Tensor,
+            inference: bool = False,
+            in_buffer: torch.Tensor = None,
+            buffer_io: Optional[ BufferIO ] = None,
             ):
-        if inference:
-            if in_buffer is None:
-                x, self.buffer = self.inference(x, self.buffer)
-                return x
-            else:
-                return self.inference(x, in_buffer)
-        else:
-            return self._forward(x)
-    
-    def inference(self, x, in_buffer):
-
-        if not self.causal:
+        if in_buffer is not None:
             raise ValueError(
                 """
-                Streaming inference is only supported for causal convolutions.
-                """
-                )
-
-        if x.shape[0] != 1:
-            raise ValueError(
-                f"""
-                Streaming inference of CausalConv1D layer only supports
-                a batch size of 1, but batch size is {x.shape[0]}.
+                The argument 'in_buffer' was removed in pytorch-tcn >= 1.2.2.
+                Instead, you should pass the input buffer as a BufferIO object
+                to the argument 'buffer_io'.
                 """
                 )
-
-        x = torch.cat(
-            (in_buffer, x),
-            -1,
-            )
-
-        out_buffer = x[:, :, -self.pad_len: ]
+        x = self.padder(x, inference=inference, buffer_io=buffer_io)
         x = super().forward(x)
-        return x, out_buffer
+        return x
+    
+    def inference(self, *args, **kwargs):
+        raise NotImplementedError(
+            """
+            The function "inference" was removed in pytorch-tcn >= 1.2.2.
+            Instead, you should use the modules forward function with the
+            argument "inference=True" enabled.
+            """
+            )
+        return
     
     def reset_buffer(self):
-        self.buffer.zero_()
-        if self.buffer.shape[2] != self.pad_len:
-            raise ValueError(
-                f"""
-                Buffer shape {self.buffer.shape} does not match the expected
-                shape (1, {self.in_channels}, {self.pad_len}).
-                """
-                )
+        self.padder.reset_buffer()
+        return
 
 
 class TemporalConvTranspose1d(nn.ConvTranspose1d):
@@ -283,11 +232,13 @@ def __init__(
                 """
                 )
 
-        # This implementation only supports kernel_size == 2 * stride with power of 2 strides
-        if kernel_size != 2 * stride or not math.log2(stride).is_integer() or stride < 2:
+        # This implementation only supports kernel_size == 2 * stride
+        if kernel_size != 2 * stride:
             raise ValueError(
                 f"""
-                This implementation only supports kernel_size == 2 * stride with power of 2 strides and stride >= 2.
+                This implementation of TemporalConvTranspose1d only
+                supports kernel_size == 2 * stride, but got 
+                kernel_size = {kernel_size} and stride = {stride}.
                 """
                 )
 
@@ -297,12 +248,12 @@ def __init__(
         self.buffer_size = (kernel_size // stride) - 1
 
         if self.causal:
-            self.pad_left = self.buffer_size
-            self.pad_right = 0
+            #self.pad_left = self.buffer_size
+            #self.pad_right = 0
             self.implicit_padding = 0
         else:
-            self.pad_left = 0
-            self.pad_right = 0
+            #self.pad_left = 0
+            #self.pad_right = 0
             self.implicit_padding = (kernel_size-stride)//2
 
         super(TemporalConvTranspose1d, self).__init__(
@@ -310,7 +261,7 @@ def __init__(
             out_channels = out_channels,
             kernel_size = kernel_size,
             stride = stride,
-            padding = self.implicit_padding, # Padding is reimplemented in this class
+            padding = self.implicit_padding, # Padding is patially reimplemented in this class
             output_padding = 0, # Output padding is not supported
             groups = groups,
             bias = bias,
@@ -320,7 +271,15 @@ def __init__(
             dtype=dtype,
             )
         
+        self.padder = TemporalPad1d(
+            padding = self.buffer_size,
+            in_channels = in_channels,
+            padding_mode = padding_mode,
+            causal = causal,
+            )
 
+        # Deprecated in pytorch-tcn >= 1.2.2
+        # Keep for backwards compatibility to load old model weights
         if buffer is None:
             buffer = torch.zeros(
                 1,
@@ -332,83 +291,44 @@ def __init__(
             buffer,
             )
 
-        if padding_mode == 'zeros':
-            self.padder = nn.ConstantPad1d(
-                    (self.pad_left, self.pad_right),
-                    0.0,
-                    )
-        elif padding_mode == 'reflect':
-            self.padder = nn.ReflectionPad1d(
-                    (self.pad_left, self.pad_right),
-                    )
-        elif padding_mode == 'replicate':
-            self.padder = nn.ReplicationPad1d(
-                    (self.pad_left, self.pad_right),
-                    )
-        elif padding_mode == 'circular':
-            self.padder = nn.CircularPad1d(
-                    (self.pad_left, self.pad_right),
-                    )
-        else:
-            raise ValueError(
-                f"""
-                padding_mode must be one of {PADDING_MODES},
-                but got {padding_mode}.
-                """
-                )
-
         return
-
-    def _forward(self, x):
-        x = self.padder(x)
-        x = super().forward(x)
-
-        if self.causal:
-            x = x[:, :, self.upsampling_factor : -self.upsampling_factor]
-
-        return x
-
-
+        
     def forward(
             self,
-            x,
-            inference=False,
-            in_buffer=None,
+            x: torch.Tensor,
+            inference: bool = False,
+            in_buffer: torch.Tensor = None,
+            buffer_io: Optional[ BufferIO ] = None,
             ):
-        if inference:
-            if in_buffer is None:
-                x, self.buffer = self.inference(x, self.buffer)
-                return x
-            else:
-                return self.inference(x, in_buffer)
-        else:
-            return self._forward(x)    
-    
-    def inference(self, x, in_buffer):
-        if not self.causal:
+        if in_buffer is not None:
             raise ValueError(
                 """
-                Streaming inference is only supported for causal convolutions.
+                The argument 'in_buffer' was removed in pytorch-tcn >= 1.2.2.
+                Instead, you should pass the input buffer as a BufferIO object
+                to the argument 'buffer_io'.
                 """
                 )
-        
-        if x.shape[0] != 1:
-            raise ValueError(
-                f"""
-                Streaming inference of CausalConv1D layer only supports
-                a batch size of 1, but batch size is {x.shape[0]}.
-                """
-                )        
-
-        x = torch.cat(
-            (in_buffer, x),
-            -1,
+        if self.causal:
+            x = self.padder(x, inference=inference, buffer_io=buffer_io)
+            x = super().forward(x)
+            x = x[:, :, self.upsampling_factor : -self.upsampling_factor]
+        else:
+            x = super().forward(x)
+            # if stride is odd, remove last element due to padding
+            if self.upsampling_factor % 2 == 1:
+                x = x[..., :-1]
+        return x
+    
+    def inference(self, *args, **kwargs):
+        raise NotImplementedError(
+            """
+            The function "inference" was removed in pytorch-tcn >= 1.2.2.
+            Instead, you should use the modules forward function with the
+            argument "inference=True" enabled.
+            """
             )
-        out_buffer = x[:, :, -self.buffer_size:]
-        x = super().forward(x)
-        x = x[:, :, self.upsampling_factor : -self.upsampling_factor]
-        return x, out_buffer
+        return
     
     def reset_buffer(self):
-        self.buffer.zero_()
+        self.padder.reset_buffer()
 
diff --git a/pytorch_tcn/pad.py b/pytorch_tcn/pad.py
new file mode 100644
index 0000000..605f6f0
--- /dev/null
+++ b/pytorch_tcn/pad.py
@@ -0,0 +1,166 @@
+import os
+import warnings
+import torch
+import torch.nn as nn
+import math
+
+from .buffer import BufferIO
+
+from typing import Optional
+from typing import Union
+from typing import List
+
+# Padding modes
+PADDING_MODES = [
+    'zeros',
+    'reflect',
+    'replicate',
+    'circular',
+]
+
+class TemporalPad1d(nn.Module):
+    def __init__(
+            self,
+            padding: int,
+            in_channels: int,
+            buffer: Optional[ Union[ float, torch.Tensor ] ] = None,
+            padding_mode: str = 'zeros',
+            causal: bool = False,
+            ):
+        super(TemporalPad1d, self).__init__()
+
+        if not isinstance(padding, int):
+            raise ValueError(
+                f"""
+                padding must be an integer, but got {type(padding)}.
+                padding must not be a tuple, because the TemporalPadding
+                will automatically determine the amount of left and right
+                padding based on the causal flag.
+                """
+                )
+
+        self.pad_len = padding
+        self.causal = causal
+
+        if causal:
+            # Padding is only on the left side
+            self.left_padding = self.pad_len
+            self.right_padding = 0
+        else:
+            # Padding is on both sides
+            self.left_padding = self.pad_len // 2
+            self.right_padding = self.pad_len - self.left_padding
+        
+        if padding_mode == 'zeros':
+            self.pad = nn.ConstantPad1d(
+                (self.left_padding, self.right_padding),
+                0.0,
+                )
+        elif padding_mode == 'reflect':
+            self.pad = nn.ReflectionPad1d(
+                (self.left_padding, self.right_padding),
+                )
+        elif padding_mode == 'replicate':
+            self.pad = nn.ReplicationPad1d(
+                (self.left_padding, self.right_padding),
+                )
+        elif padding_mode == 'circular':
+            self.pad = nn.CircularPad1d(
+                (self.left_padding, self.right_padding),
+                )
+        else:
+            raise ValueError(
+                f"""
+                padding_mode must be one of {PADDING_MODES},
+                but got {padding_mode}.
+                """
+                )
+        
+        # Buffer is used for streaming inference
+        if buffer is None:
+            buffer = torch.zeros(
+                1,
+                in_channels,
+                self.pad_len,
+                )
+        elif isinstance(buffer, (int, float)):
+            buffer = torch.full(
+                size = (1, in_channels, self.pad_len),
+                fill_value = buffer,
+            )
+        elif not isinstance(buffer, torch.Tensor):
+            raise ValueError(
+                f"""
+                The argument 'buffer' must be None or of type float,
+                int, or torch.Tensor, but got {type(buffer)}.
+                """
+                )
+        
+        # Register buffer as a persistent buffer which is available as self.buffer
+        self.register_buffer(
+            'buffer',
+            buffer,
+            )
+        
+        return
+    
+    def pad_inference(
+            self,
+            x: torch.Tensor,
+            buffer_io: Optional[ BufferIO ] = None,
+            ):
+
+        if not self.causal:
+            raise ValueError(
+                """
+                Streaming inference is only supported for causal convolutions.
+                """
+                )
+
+        if x.shape[0] != 1:
+            raise ValueError(
+                f"""
+                Streaming inference requires a batch size
+                of 1, but batch size is {x.shape[0]}.
+                """
+                )
+        
+        if buffer_io is None:
+            in_buffer = self.buffer
+        else:
+            in_buffer = buffer_io.next_in_buffer()
+
+        x = torch.cat(
+            (in_buffer, x),
+            -1,
+            )
+
+        out_buffer = x[:, :, -self.pad_len: ]
+        if buffer_io is None:
+            self.buffer = out_buffer
+        else:
+            buffer_io.append_out_buffer(out_buffer)
+
+        return x
+    
+    def forward(
+            self,
+            x: torch.Tensor,
+            inference: bool = False,
+            buffer_io: Optional[ BufferIO ] = None,
+            ):
+        if inference:
+            x = self.pad_inference(x, buffer_io=buffer_io)
+        else:
+            x = self.pad(x)
+        return x
+    
+    def reset_buffer(self):
+        self.buffer.zero_()
+        if self.buffer.shape[2] != self.pad_len:
+            raise ValueError(
+                f"""
+                Buffer shape {self.buffer.shape} does not match the expected
+                shape (1, {self.in_channels}, {self.pad_len}).
+                """
+                )
\ No newline at end of file
diff --git a/pytorch_tcn/tcn.py b/pytorch_tcn/tcn.py
index b8ed821..ec7fda8 100644
--- a/pytorch_tcn/tcn.py
+++ b/pytorch_tcn/tcn.py
@@ -61,14 +61,24 @@ def _check_activation_arg(
                 torch.nn.Module object as the 'activation' argument.
                 """
                 )
-    elif not isinstance( activation, nn.Module ):
-        raise ValueError(
-            f"""
-            The argument '{arg_name}' must either be a valid string or
-            a torch.nn.Module object, but {activation} was passed,
-            which is of type {type(activation)}.
-            """
-            )
+    else:
+        try:
+            if not isinstance( activation(), nn.Module ):
+                raise ValueError(
+                    f"""
+                    The argument '{arg_name}' must either be a valid string or
+                    a torch.nn.Module object, but an object of type {type(activation())}
+                    was passed.
+                    """
+                    )
+        except:
+            raise ValueError(
+                f"""
+                The argument '{arg_name}' must either be a valid string or
+                a torch.nn.Module object, but an object of type {type(activation)}
+                was passed.
+                """
+                )
     return
 
 def _check_generic_input_arg(
@@ -89,9 +99,12 @@ def get_kernel_init_fn(
         name: str,
         activation: str,
         ) -> Tuple[ nn.Module, dict ]:
-    if isinstance( activation, nn.Module ):
-        return kernel_init_fn[ name ], dict()
-    # TODO: this means no gain is used for custom activation functions
+    try:
+        if isinstance( activation(), nn.Module ):
+            return kernel_init_fn[ name ], dict()
+        # TODO: this means no gain is used for custom activation functions
+    except:
+        pass
         
     if name not in kernel_init_fn.keys():
         raise ValueError(


Test patch:
diff --git a/tests/unit/test_conv.py b/tests/unit/test_conv.py
index 528e8ea..f0063f1 100644
--- a/tests/unit/test_conv.py
+++ b/tests/unit/test_conv.py
@@ -1,150 +1,256 @@
+
+import os
+import tempfile
 import torch
+import torch.nn as nn
 import unittest
 
-import torch.nn as nn
 from pytorch_tcn.conv import TemporalConv1d, TemporalConvTranspose1d
-import tempfile
-import os
+from pytorch_tcn.buffer import BufferIO
+
+
+class ConvModel(nn.Module):
+            def __init__(
+                    self,
+                    layer_1,
+                    layer_2,
+                    ):
+                super(ConvModel, self).__init__()
+                self.conv1 = layer_1
+                self.conv2 = layer_2
+                return
+
+            def forward(self, x, in_buffers):
+                buffer_io = BufferIO( in_buffers = in_buffers )
+                with torch.no_grad():
+                    x = self.conv1(
+                        x=x,
+                        inference=True,
+                        buffer_io=buffer_io,
+                    )
+                    x = self.conv2(
+                        x=x,
+                        inference=True,
+                        buffer_io=buffer_io,
+                    )
+                out_buffers = buffer_io.out_buffers
+                return x, out_buffers
 
 class ConvolutionTest(unittest.TestCase):
-    def test_causal_internal_buffer_convolution(self):
-        # Define your input tensor
-        input_tensor = torch.randn(1, 3, 32)  # Example input tensor shape: (batch_size, channels, time_steps)
+        
+    def test_conv1d_streaming_with_internal_buffer(self):
+        # Example input tensor shape: (batch_size, channels, time_steps)
+        batch_size = 1
+        in_channels = 3
+        time_steps = 32
+        input_tensor = torch.randn(batch_size, in_channels, time_steps)
 
-        # Define your convolutional layer
-        in_channels = 3  # Number of input channels
+        # Define the convolutional layer
         out_channels = 16  # Number of output channels
         kernel_size = 3  # Size of the convolutional kernel
-        stride = 1  # Stride value for the convolution
-        dilation = 1  # Dilation value for the convolution
         conv_layer = TemporalConv1d(
             in_channels = in_channels,
             out_channels = out_channels,
             kernel_size = kernel_size,
-            stride = stride,
-            dilation = dilation,
+            stride = 1,
+            dilation = 1,
             causal=True,
             )
-
-        self.assertEqual(conv_layer.buffer.shape, (1, in_channels, kernel_size - 1))  # Example assertion for buffer shape
+        
+        # Test the buffer shape
+        self.assertEqual(
+            conv_layer.buffer.shape,
+            (1, in_channels, kernel_size - 1),
+            )
 
         with torch.no_grad():
 
             # Apply the convolutional layer to the input tensor
             output_tensor = conv_layer(input_tensor)
 
-            # Perform your assertions or checks here
-            self.assertEqual(output_tensor.shape, (1, 16, 32))  # Example assertion for output tensor shape
+            # Test the output tensor shape
+            self.assertEqual(
+                output_tensor.shape,
+                (batch_size, out_channels, time_steps),
+                )
 
             # Compare with running inference in a loop, with internal buffers
             conv_layer.reset_buffer()
             for t in range(0,input_tensor.shape[2]-1):
                 input_slice = input_tensor[:, :, t:t+1]
                 
-                output_slice = conv_layer.forward(input_slice, inference=True)
+                output_slice = conv_layer(
+                    input_slice,
+                    inference=True,
+                    )
                 reference_output_slice = output_tensor[:,:,t:t+1]
 
                 assert(torch.allclose(output_slice, reference_output_slice, atol=1e-5))
 
+        return
 
 
-    def test_causal_external_buffer_convolution(self):
-        # Define your input tensor
-        input_tensor = torch.randn(1, 3, 32)  # Example input tensor shape: (batch_size, channels, time_steps)
 
-        # Define your convolutional layer
-        in_channels = 3  # Number of input channels
+    def test_conv1d_streaming_with_external_buffer(self):
+        # Example input tensor shape: (batch_size, channels, time_steps)
+        batch_size = 1
+        in_channels = 3
+        time_steps = 32
+        input_tensor = torch.randn(batch_size, in_channels, time_steps)
+
+        # Define the convolutional layer
         out_channels = 16  # Number of output channels
         kernel_size = 3  # Size of the convolutional kernel
-        stride = 1  # Stride value for the convolution
-        dilation = 1  # Dilation value for the convolution
         conv_layer = TemporalConv1d(
             in_channels = in_channels,
             out_channels = out_channels,
             kernel_size = kernel_size,
-            stride = stride,
-            dilation = dilation,
+            stride = 1,
+            dilation = 1,
             causal=True,
             )
-
-        self.assertEqual(conv_layer.buffer.shape, (1, in_channels, kernel_size - 1))  # Example assertion for buffer shape
+        
+        # Test the buffer shape
+        self.assertEqual(
+            conv_layer.buffer.shape,
+            (1, in_channels, kernel_size - 1),
+            )
 
         with torch.no_grad():
 
             # Apply the convolutional layer to the input tensor
             output_tensor = conv_layer(input_tensor)
 
-            # Perform your assertions or checks here
-            self.assertEqual(output_tensor.shape, (1, 16, 32))  # Example assertion for output tensor shape
+            # Test the output tensor shape
+            self.assertEqual(
+                output_tensor.shape,
+                (batch_size, out_channels, time_steps),
+                )
+
+            # Define the BufferIO object
             in_buffer = torch.zeros(1, in_channels, kernel_size - 1)
+            buffer_io = BufferIO( in_buffers = [in_buffer] )
 
             # Compare with running inference in a loop with external buffers
             for t in range(0,input_tensor.shape[2]):
                 input_slice = input_tensor[:, :, t:t+1]
                 
-                output_slice, out_buffer = conv_layer.inference(input_slice, in_buffer=in_buffer)
+                output_slice = conv_layer(
+                    input_slice,
+                    inference=True,
+                    buffer_io=buffer_io,
+                    )
                 reference_output_slice = output_tensor[:,:,t:t+1]
 
                 assert(torch.allclose(output_slice, reference_output_slice, atol=1e-5))
 
-                in_buffer = out_buffer
+                buffer_io.step()
+        return
 
-    def test_causal_onnx_convolution(self):
+    def test_conv1d_streaming_with_onnx(self):
         import onnxruntime as ort
 
-        # Define your input tensor
-        input_tensor = torch.randn(1, 3, 32)  # Example input tensor shape: (batch_size, channels, time_steps)
+        # Example input tensor shape: (batch_size, channels, time_steps)
+        batch_size = 1
+        in_channels = 3
+        time_steps = 32
+        input_tensor = torch.randn(batch_size, in_channels, time_steps)
 
-        # Define your convolutional layer
-        in_channels = 3  # Number of input channels
-        out_channels = 16  # Number of output channels
+        # Define the convolutional layer
+        hidden_channels = 4
+        out_channels = 2  # Number of output channels
         kernel_size = 3  # Size of the convolutional kernel
-        stride = 1  # Stride value for the convolution
-        dilation = 1  # Dilation value for the convolution
-        conv_layer = TemporalConv1d(
+        conv_layer_1 = TemporalConv1d(
             in_channels = in_channels,
+            out_channels = hidden_channels,
+            kernel_size = kernel_size,
+            stride = 1,
+            dilation = 1,
+            causal=True,
+            )
+        conv_layer_2 = TemporalConv1d(
+            in_channels = hidden_channels,
             out_channels = out_channels,
             kernel_size = kernel_size,
-            stride = stride,
-            dilation = dilation,
+            stride = 1,
+            dilation = 1,
             causal=True,
             )
 
-        class ConvModel(nn.Module):
-            def __init__(self, conv_layer):
-                super(ConvModel, self).__init__()
-                self.conv_layer = conv_layer
+        # Define model
+        model = ConvModel(
+            layer_1=conv_layer_1,
+            layer_2=conv_layer_2,
+            )
+        
+        # Define the BufferIO object
+        in_buffer_1 = torch.zeros(1, in_channels, kernel_size - 1)
+        in_buffer_2 = torch.zeros(1, hidden_channels, kernel_size - 1)
+        in_buffers = [in_buffer_1, in_buffer_2]
 
-            def forward(self, x, in_buffer):
-                with torch.no_grad():
-                    return self.conv_layer.inference(x, in_buffer)
+        # Reference output (full tensor)
+        out_tensor, _ = model(input_tensor, in_buffers)
 
-        model = ConvModel(conv_layer)
-        in_buffer = torch.zeros(1, in_channels, kernel_size - 1)
+        # Test the streaming inference with ONNX
         input_slice = input_tensor[:, :, 0:1]
 
+        onnx_stream = []
         with tempfile.TemporaryDirectory() as temp_dir:
             onnx_model_name = os.path.join(temp_dir, "test_conv_model.onnx")
-            torch.onnx.export(model, (input_slice, in_buffer), onnx_model_name, 
-                              input_names=['in_x', 'in_buffer'], output_names=['out_x', 'out_buffer'], 
-                              opset_version=9, export_params=True)
-
-
+            torch.onnx.export(
+                model=model,
+                args=(input_slice, in_buffers),
+                f=onnx_model_name, 
+                input_names=['in_x', 'in_buffer_1', 'in_buffer_2'],
+                output_names=['out_x', 'out_buffer_1', 'out_buffer_2'],
+                opset_version=9,
+                export_params=True,
+                )
+
+            # ONNX streaming output
             ort_session = ort.InferenceSession(onnx_model_name)
-
-            for t in range(0,input_tensor.shape[2]-1):
-                input_slice = input_tensor[:, :, t:t+1]
+            for t in range(0,input_tensor.shape[2]):
+                input_slice = input_tensor[..., t:t+1]
                 
-                output_slice, out_buffer = model(input_slice, in_buffer)
-                onnx_outputs = ort_session.run(None, {'in_x': input_slice.numpy(), 'in_buffer': in_buffer.numpy()})
-
-                assert(torch.allclose(output_slice, torch.tensor(onnx_outputs[0]), atol=1e-5))
-                assert(torch.allclose(out_buffer, torch.tensor(onnx_outputs[1]), atol=1e-5))
+                # Reference output (streamed)
+                out_tensor_stream, out_buffers_stream = model(
+                    input_slice,
+                    in_buffers,
+                    )
+                onnx_outputs = ort_session.run(
+                    None,
+                    {
+                        'in_x': input_slice.numpy(),
+                        'in_buffer_1': in_buffers[0].numpy(),
+                        'in_buffer_2': in_buffers[1].numpy(),
+                        }
+                    )
+                onnx_output_slice = torch.tensor(onnx_outputs[0])
+                onnx_out_buffers = [ torch.tensor(b) for b in onnx_outputs[1:] ]
+                onnx_stream.append(onnx_output_slice)
+
+                # Test that onnx matches reference (streamed inference)
+                #   - model output:
+                assert(torch.allclose(out_tensor_stream, onnx_output_slice, atol=1e-5))
+                #   - buffer output:
+                for out_b, onnx_out_b in zip(out_buffers_stream, onnx_out_buffers):
+                    assert(torch.allclose(out_b, onnx_out_b, atol=1e-5))
+
+                # Move buffers
+                in_buffers = onnx_out_buffers
+
+        # Test that onnx matches reference (full inference)
+        assert(torch.allclose(out_tensor, torch.cat(onnx_stream, dim=2), atol=1e-5))
+            
+        return
         
 
     def test_noncausal_convolution(self):
-        # Define your input tensor
-        input_tensor = torch.randn(1, 3, 32)  # Example input tensor shape: (batch_size, channels, time_steps)
+        # Example input tensor shape: (batch_size, channels, time_steps)
+        batch_size = 1
+        in_channels = 3
+        time_steps = 32
+        input_tensor = torch.randn(batch_size, in_channels, time_steps)
 
         # Define your convolutional layer
         in_channels = 3  # Number of input channels
@@ -184,16 +290,17 @@ def test_noncausal_convolution(self):
             assert(torch.allclose(output_tensor, reference_output_tensor, atol=1e-5))
 
 
-    def test_transpose_causal_internal_buffer_convolution(self):
+    def test_convtranspose1d_streaming_with_internal_buffers(self):
         # Define your input tensor
         input_tensor = torch.randn(1, 3, 32)  # Example input tensor shape: (batch_size, channels, time_steps)
 
         # Define your convolutional layer
         in_channels = 3  # Number of input channels
         out_channels = 16  # Number of output channels
-        kernel_size = 8  # Size of the convolutional kernel
         stride = 4  # Stride value for the convolution
-        conv_layer = TemporalConvTranspose1d(
+        kernel_size = 2 * stride
+
+        conv_t = TemporalConvTranspose1d(
             in_channels = in_channels,
             out_channels = out_channels,
             kernel_size = kernel_size,
@@ -204,19 +311,19 @@ def test_transpose_causal_internal_buffer_convolution(self):
         with torch.no_grad():
 
             # Apply the convolutional layer to the input tensor
-            output_tensor = conv_layer(input_tensor, inference=True)
+            output_tensor = conv_t(input_tensor)
 
             # Compare with running inference in a loop, with internal buffers
-            conv_layer.reset_buffer()
+            conv_t.reset_buffer()
             for t in range(0,input_tensor.shape[2]):
                 input_slice = input_tensor[:, :, t:t+1]
                 
-                output_slice = conv_layer(input_slice, inference=True)
+                output_slice = conv_t(input_slice, inference=True)
                 reference_output_slice = output_tensor[:,:,t*stride:(t+1)*stride]
 
                 assert(torch.allclose(output_slice, reference_output_slice, atol=1e-5))
 
-    def test_transpose_causal_external_buffer_convolution(self):
+    def test_convtranspose1d_streaming_with_external_buffers(self):
         # Define your input tensor
         input_tensor = torch.randn(1, 3, 32)
 
@@ -225,7 +332,7 @@ def test_transpose_causal_external_buffer_convolution(self):
         out_channels = 16
         kernel_size = 8
         stride = 4
-        conv_layer = TemporalConvTranspose1d(
+        conv_t = TemporalConvTranspose1d(
             in_channels = in_channels,
             out_channels = out_channels,
             kernel_size = kernel_size,
@@ -236,22 +343,27 @@ def test_transpose_causal_external_buffer_convolution(self):
         with torch.no_grad():
                 
                 # Apply the convolutional layer to the input tensor
-                output_tensor = conv_layer(input_tensor, inference=True)
+                output_tensor = conv_t(input_tensor)
     
-                in_buffer = torch.zeros(1, in_channels, conv_layer.buffer_size)
+                in_buffer = torch.zeros(1, in_channels, conv_t.buffer_size)
+                buffer_io = BufferIO( in_buffers = [in_buffer] )
     
                 # Compare with running inference in a loop with external buffers
                 for t in range(0,input_tensor.shape[2]):
                     input_slice = input_tensor[:, :, t:t+1]
                     
-                    output_slice, out_buffer = conv_layer.inference(input_slice, in_buffer=in_buffer)
+                    output_slice = conv_t(
+                        input_slice,
+                        inference=True,
+                        buffer_io=buffer_io,
+                        )
                     reference_output_slice = output_tensor[:,:,t*stride:(t+1)*stride]
     
                     assert(torch.allclose(output_slice, reference_output_slice, atol=1e-5))
     
-                    in_buffer = out_buffer
-          
-    def test_transpose_causal_onnx_convolution(self):
+                    buffer_io.step()
+
+    def test_convtranspose1d_streaming_with_onnx(self):
         import onnxruntime as ort
 
         # Define your input tensor
@@ -259,47 +371,92 @@ def test_transpose_causal_onnx_convolution(self):
 
         # Define your convolutional layer
         in_channels = 3
+        hidden_channels = 4
         out_channels = 16
         kernel_size = 8
         stride = 4
-        conv_layer = TemporalConvTranspose1d(
+        conv_layer_1 = TemporalConvTranspose1d(
             in_channels = in_channels,
-            out_channels = out_channels,
+            out_channels = hidden_channels,
             kernel_size = kernel_size,
             stride = stride,
             causal=True,
             )
+        conv_layer_2 = TemporalConvTranspose1d(
+            in_channels = hidden_channels,
+            out_channels = out_channels,
+            kernel_size = 8,
+            stride = 4,
+            causal=True,
+            )
+                
+        # Define model
+        model = ConvModel(
+            layer_1=conv_layer_1,
+            layer_2=conv_layer_2,
+            )
+        
+        # Define the BufferIO object
+        in_buffer_1 = torch.zeros(1, in_channels, conv_layer_1.buffer_size)
+        in_buffer_2 = torch.zeros(1, hidden_channels, conv_layer_2.buffer_size)
+        in_buffers = [in_buffer_1, in_buffer_2]
 
-        class ConvModel(nn.Module):
-            def __init__(self, conv_layer):
-                super(ConvModel, self).__init__()
-                self.conv_layer = conv_layer
+        # Reference output (full tensor)
+        out_tensor, _ = model(input_tensor, in_buffers)
 
-            def forward(self, x, in_buffer):
-                with torch.no_grad():
-                    return self.conv_layer.inference(x, in_buffer)
-                
-        model = ConvModel(conv_layer)
-        in_buffer = torch.zeros(1, in_channels, conv_layer.buffer_size)
+        # Test the streaming inference with ONNX
         input_slice = input_tensor[:, :, 0:1]
 
+        onnx_stream = []
         with tempfile.TemporaryDirectory() as temp_dir:
             onnx_model_name = os.path.join(temp_dir, "test_conv_model.onnx")
-            torch.onnx.export(model, (input_slice, in_buffer), onnx_model_name, 
-                              input_names=['in_x', 'in_buffer'], output_names=['out_x', 'out_buffer'], 
-                              opset_version=9, export_params=True)
-
-
+            torch.onnx.export(
+                model=model,
+                args=(input_slice, in_buffers),
+                f=onnx_model_name, 
+                input_names=['in_x', 'in_buffer_1', 'in_buffer_2'],
+                output_names=['out_x', 'out_buffer_1', 'out_buffer_2'],
+                opset_version=9,
+                export_params=True,
+                )
+
+            # ONNX streaming output
             ort_session = ort.InferenceSession(onnx_model_name)
-
             for t in range(0,input_tensor.shape[2]):
-                input_slice = input_tensor[:, :, t:t+1]
+                input_slice = input_tensor[..., t:t+1]
                 
-                output_slice, out_buffer = model(input_slice, in_buffer)
-                onnx_outputs = ort_session.run(None, {'in_x': input_slice.numpy(), 'in_buffer': in_buffer.numpy()})
-
-                assert(torch.allclose(output_slice, torch.tensor(onnx_outputs[0]), atol=1e-5))
-                assert(torch.allclose(out_buffer, torch.tensor(onnx_outputs[1]), atol=1e-5))
+                # Reference output (streamed)
+                out_tensor_stream, out_buffers_stream = model(
+                    input_slice,
+                    in_buffers,
+                    )
+                onnx_outputs = ort_session.run(
+                    None,
+                    {
+                        'in_x': input_slice.numpy(),
+                        'in_buffer_1': in_buffers[0].numpy(),
+                        'in_buffer_2': in_buffers[1].numpy(),
+                        }
+                    )
+                onnx_output_slice = torch.tensor(onnx_outputs[0])
+                onnx_out_buffers = [ torch.tensor(b) for b in onnx_outputs[1:] ]
+                onnx_stream.append(onnx_output_slice)
+
+                # Test that onnx matches reference (streamed inference)
+                #   - model output:
+                assert(torch.allclose(out_tensor_stream, onnx_output_slice, atol=1e-5))
+                #   - buffer output:
+                for out_b, onnx_out_b in zip(out_buffers_stream, onnx_out_buffers):
+                    assert(torch.allclose(out_b, onnx_out_b, atol=1e-5))
+
+                # Move buffers
+                in_buffers = onnx_out_buffers
+
+
+        # Test that onnx matches reference (full inference)
+        assert(torch.allclose(out_tensor, torch.cat(onnx_stream, dim=2), atol=1e-5))
+            
+        return
 
     def test_transpose_noncausal_convolution(self):
         # Define your input tensor
diff --git a/tests/unit/test_tcn.py b/tests/unit/test_tcn.py
index 07de96e..e4a4f28 100644
--- a/tests/unit/test_tcn.py
+++ b/tests/unit/test_tcn.py
@@ -113,6 +113,13 @@ def __init__(self, methodName: str = "runTest") -> None:
                     ),
                 expected_error = None,
             ),
+            dict(
+                kwargs = dict(
+                    activation = [torch.nn.SiLU],
+                    use_gate = [False],
+                    ),
+                expected_error = None,
+            ),
             dict(
                 kwargs = dict( activation = [ 'invalid' ] ),
                 expected_error = ValueError,