[BugFix] ConvNet forward method with tensors of more than 4 dimensions (#686)

* cnn forward fix

* more general code

* cnn testing

* precommit run check

* convnet tests

Author: albertbou92

test/test_modules.py

@@ -124,7 +124,9 @@ def test_mlp(
 )
 @pytest.mark.parametrize("squeeze_output", [False])
 @pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("batch", [(2,), (2, 2)])
 def test_convnet(
+    batch,
     in_features,
     depth,
     num_cells,
@@ -145,7 +147,6 @@ def test_convnet(
     seed=0,
 ):
     torch.manual_seed(seed)
-    batch = 2
     convnet = ConvNet(
         in_features=in_features,
         depth=depth,
@@ -165,9 +166,9 @@ def test_convnet(
     )
     if in_features is None:
         in_features = 5
-    x = torch.randn(batch, in_features, input_size, input_size, device=device)
+    x = torch.randn(*batch, in_features, input_size, input_size, device=device)
     y = convnet(x)
-    assert y.shape == torch.Size([batch, expected_features])
+    assert y.shape == torch.Size([*batch, expected_features])
 
 
 @pytest.mark.parametrize(

torchrl/modules/models/models.py

@@ -463,6 +463,15 @@ def _make_net(self, device: Optional[DEVICE_TYPING]) -> nn.Module:
             layers.append(Squeeze2dLayer())
         return layers
 
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        *batch, C, L, W = inputs.shape
+        if len(batch) > 1:
+            inputs = inputs.flatten(0, len(batch) - 1)
+        out = super(ConvNet, self).forward(inputs)
+        if len(batch) > 1:
+            out = out.unflatten(0, batch)
+        return out
+
 
 class DuelingMlpDQNet(nn.Module):
     """Creates a Dueling MLP Q-network.