Fix #4

Author: 1adrianb

bnn/layers/conv.py

@@ -36,6 +36,9 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         input_proc = self.activation_pre_process(input)
         input_proc = self._conv_forward(input_proc, self.weight_pre_process(self.weight), bias=self.bias)
 
+        if isinstance(input_proc, tuple) and len(input_proc) == 1:
+            input_proc = input_proc[0]
+        
         return self.activation_post_process(
             input_proc,
             input
@@ -91,6 +94,9 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         input_proc = self.activation_pre_process(input)
         input_proc = self._conv_forward(input_proc, self.weight_pre_process(self.weight), bias=self.bias)
 
+        if isinstance(input_proc, tuple) and len(input_proc) == 1:
+            input_proc = input_proc[0]
+        
         return self.activation_post_process(
             input_proc,
             input