From 4dfb0b1f354ca4c700e584e4c1ea68ded84442d7 Mon Sep 17 00:00:00 2001 From: Francois Drielsma Date: Thu, 9 Nov 2023 15:17:31 -0800 Subject: [PATCH] Fixed batch size problem in unwrapper when there's multiple volumes --- mlreco/trainval.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlreco/trainval.py b/mlreco/trainval.py index c6bd9dfc..fc210bb8 100644 --- 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: