diff --git a/flood_forecast/pytorch_training.py b/flood_forecast/pytorch_training.py
index d0261d1bd..a22cf86a4 100644
--- a/flood_forecast/pytorch_training.py
+++ b/flood_forecast/pytorch_training.py
@@ -231,15 +231,10 @@ def compute_loss(labels, output, src, criterion, validation_dataset, probabilist
             output = validation_dataset.inverse_scale(output.cpu())
             labels = validation_dataset.inverse_scale(labels.cpu())
         elif len(output.shape) == 3:
-            print("original output shape ")
-            print(output.shape)
             output = output.cpu().numpy().transpose(0, 2, 1)
             labels = labels.cpu().numpy().transpose(0, 2, 1)
             output = validation_dataset.inverse_scale(torch.from_numpy(output))
             labels = validation_dataset.inverse_scale(torch.from_numpy(labels))
-            print("Output shape is ")
-            print(output)
-            print(labels)
             stuff = src.cpu().numpy().transpose(0, 2, 1)
             src = validation_dataset.inverse_scale(torch.from_numpy(stuff))
         else:
@@ -255,8 +250,6 @@ def compute_loss(labels, output, src, criterion, validation_dataset, probabilist
         assert len(labels.shape) == len(output.shape)
         loss = criterion(labels.float(), output, src, m)
     else:
-        print(output.shape)
-        print(labels.shape)
         assert len(labels.shape) == len(output.shape)
         assert labels.shape[0] == output.shape[0]
         loss = criterion(output, labels.float())
diff --git a/tests/transformer_bottleneck.json b/tests/transformer_bottleneck.json
index e2e1761d5..1f4482155 100644
--- a/tests/transformer_bottleneck.json
+++ b/tests/transformer_bottleneck.json
@@ -21,7 +21,7 @@
        "test_path": "tests/test_data/keag_small.csv",
        "batch_size":4,
        "forecast_history":5,
-       "forecast_length":5,
+       "forecast_length":1,
        "train_end": 100,
        "valid_start":301,
        "valid_end": 401,
@@ -67,7 +67,7 @@
           "dataset_params":{
              "file_path": "tests/test_data/keag_small.csv",
              "forecast_history":5,
-             "forecast_length":5,
+             "forecast_length":1,
              "relevant_cols": ["cfs", "precip", "temp"],
              "target_col": ["cfs"],
              "scaling": "RobustScaler",