fix: add dummy batch dimension to meet the requirements of the batch …

…norm
es-ude · Aug 12, 2023 · 0f499f0 · 0f499f0
1 parent 7642f80
commit 0f499f0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/elasticai/creator/nn/batch_normed_conv1d/layer.py b/elasticai/creator/nn/batch_normed_conv1d/layer.py
@@ -58,15 +58,16 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         input_shape = (
             (inputs.shape[0], self._conv1d.in_channels, -1)
             if has_batches
-            else (self._conv1d.in_channels, -1)
+            else (1, self._conv1d.in_channels, -1)
         )
         output_shape = (inputs.shape[0], -1) if has_batches else (-1,)
+
         x = inputs.view(*input_shape)
         x = self._conv1d(x)
         x = self._batch_norm(x)
         x = self._arithmetics.quantize(x)
-        outputs = x.view(*output_shape)
-        return outputs
+
+        return x.view(*output_shape)
 
     def translate(self, name: str) -> FPConv1dDesign:
         def float_to_signed_int(value: float | list) -> int | list:

diff --git a/elasticai/creator/nn/batch_normed_linear/layer.py b/elasticai/creator/nn/batch_normed_linear/layer.py
@@ -47,9 +47,16 @@ def __init__(
         )
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
-        x = self._linear(inputs)
+        has_batches = inputs.dim() == 2
+        input_shape = inputs.shape if has_batches else (1, -1)
+        output_shape = (inputs.shape[0], -1) if has_batches else (-1,)
+
+        x = inputs.view(*input_shape)
+        x = self._linear(x)
         x = self._batch_norm(x)
-        return self._arithmetics.quantize(x)
+        x = self._arithmetics.quantize(x)
+
+        return x.view(*output_shape)
 
     def translate(self, name: str) -> FPLinearDesign:
         def float_to_signed_int(value: float | list) -> int | list: