fix: make two-stage extractor work with empty batches (#56)

tilman151 · Jan 26, 2024 · 5094a82 · 5094a82
1 parent 15f0cc8
commit 5094a82
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/rul_adapt/model/two_stage.py b/rul_adapt/model/two_stage.py
@@ -52,7 +52,8 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         batch_size, upper_seq_len, input_channels, lower_seq_len = inputs.shape
         inputs = inputs.reshape(-1, input_channels, lower_seq_len)
         inputs = self.lower_stage(inputs)
-        inputs = inputs.reshape(batch_size, upper_seq_len, -1)
+        _, lower_output_units = inputs.shape
+        inputs = inputs.reshape(batch_size, upper_seq_len, lower_output_units)
         inputs = torch.transpose(inputs, 1, 2)
         inputs = self.upper_stage(inputs)
 

diff --git a/tests/test_model/test_two_stage.py b/tests/test_model/test_two_stage.py
@@ -44,3 +44,7 @@ def test_forward_upper_lower_interaction(inputs, extractor):
     outputs = extractor(inputs)
 
     assert torch.allclose(upper_outputs, outputs[3])
+
+
+def test_forward_with_empty_input(extractor):
+    output = extractor(torch.empty(0, 4, 3, 64))