let user decide evaluation batch size

src/configs/experiment/navier_stokes.yaml

@@ -10,7 +10,7 @@ name: "NavierStokes"
 
 datamodule:
   batch_size: 32
-  eval_batch_size: ${datamodule.batch_size}
+  eval_batch_size: 4     # effectively eval_batch_size *= number of predictions ($module.num_predictions; default=20)
   horizon: 16
   prediction_horizon: 64
   window: 1

src/configs/experiment/oisst_pacific.yaml

@@ -9,7 +9,7 @@ name: "OISSTv2PacificSubset"
 
 datamodule:
   batch_size: 64
-  eval_batch_size: ${datamodule.batch_size}
+  eval_batch_size: 6   # effectively eval_batch_size *= number of predictions ($module.num_predictions; default=20)
   horizon: 7
   prediction_horizon: 7
   boxes: [84, 85, 86, 87, 88, 89, 108, 109, 110, 111, 112]

src/configs/experiment/spring_mesh.yaml

@@ -11,7 +11,7 @@ name: "SpringMesh"
 datamodule:
   physical_system: 'spring-mesh'
   batch_size: 64
-  eval_batch_size: 512
+  eval_batch_size: 16  # effectively eval_batch_size *= number of predictions ($module.num_predictions; default=20)
   horizon: 134
   prediction_horizon: 804
   window: 1

src/utilities/config_utils.py

@@ -1,4 +1,3 @@
-import math
 import os
 import time
 import warnings
@@ -464,13 +463,12 @@ def check_config_values(config: DictConfig):
             config.save_config_to_wandb = False
 
         if config.module.get("num_predictions", 1) > 1:
-            # adapt the evaluation batch size to the number of predictions
             bs, ebs = config.datamodule.batch_size, config.datamodule.eval_batch_size
             if ebs >= bs:
-                # reduce the eval batch size to account for the number of predictions
-                config.datamodule.eval_batch_size = max(1, int(bs // math.sqrt(config.module.num_predictions)))
+                effective_ebs = ebs * config.module.num_predictions
                 log.info(
-                    f"Reducing eval batch size from {ebs} to {config.datamodule.eval_batch_size} to match the number of predictions."
+                    f"Note that the effective evaluation batch size will be multiplied by the number of "
+                    f"predictions={config.module.num_predictions} for a total of {effective_ebs}!"
                 )