PT forward, better error on missing batch dim

#1385
rwth-i6 · Aug 29, 2023 · 6839833 · 6839833
1 parent 5e8e145
commit 6839833
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/returnn/torch/engine.py b/returnn/torch/engine.py
@@ -677,6 +677,8 @@ def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIf
         batch_dim = _get_batch_dim_from_extern_data(self.extern_data)
 
         def _get_tensor_wo_batch_numpy(x: Tensor) -> Tensor:
+            if batch_dim not in x.dims:
+                raise Exception(f"Expected {batch_dim} in {x}.")
             if x.dims.index(batch_dim) != 0:
                 x = x.copy_move_axis(x.dims.index(batch_dim), 0)