Fixed batch size problem in unwrapper when there's multiple volumes

DeepLearnPhysics · Nov 9, 2023 · 4dfb0b1 · 4dfb0b1
1 parent 92495b5
commit 4dfb0b1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mlreco/trainval.py b/mlreco/trainval.py
@@ -227,7 +227,7 @@ def forward(self, data_iter, iteration=None):
 
             # Unwrap output, if requested
             if unwrap:
-                unwrapper.batch_size = len(input_data['index'][0]) * self._num_volumes
+                unwrapper.batch_size = len(input_data['index'][0])
                 input_data, res = unwrapper(input_data, res)
             else:
                 if 'index' in input_data: