add shape test

Author: phoeenniixx

pytorch_forecasting/tests/test_all_estimators_v2.py

@@ -6,6 +6,7 @@
 import lightning.pytorch as pl
 from lightning.pytorch.callbacks import EarlyStopping
 from lightning.pytorch.loggers import TensorBoardLogger
+import torch
 import torch.nn as nn
 
 from pytorch_forecasting.tests.test_all_estimators import (
@@ -77,6 +78,19 @@ def _integration(
     test_outputs = trainer.test(net, dataloaders=test_dataloader)
     assert len(test_outputs) > 0
 
+    # todo: add the predict pipeline and make this test cleaner
+    x, y = next(iter(test_dataloader))
+    net.eval()
+    with torch.no_grad():
+        output = net(x)
+    net.train()
+    prediction = output["prediction"]
+    n_dims = len(prediction.shape)
+    assert n_dims == 3, (
+        f"Prediction output must be 3D, but got {n_dims}D tensor "
+        f"with shape {output.shape}"
+    )
+
     shutil.rmtree(tmp_path, ignore_errors=True)