dims check

Co-authored-by: Vincent Moens <[email protected]>

torchrl/modules/models/models.py

@@ -695,7 +695,10 @@ def _make_net(self, device: Optional[DEVICE_TYPING]) -> nn.Module:
  return layers
 
  def forward(self, inputs: torch.Tensor) -> torch.Tensor:
- *batch, C, D, L, W = inputs.shape
+ try:
+ *batch, C, D, L, W = inputs.shape
+ except ValueError as err:
+ raise ValueError(f"The input value of {self.__class__.__name__} must have at least 4 dimensions, got {inputs.ndim} instead.") from err
  if len(batch) > 1:
  inputs = inputs.flatten(0, len(batch) - 1)
  out = super().forward(inputs)