fixed aggregation when 0-dim torch tensors on gpu are passed

analysiscenter · Aug 1, 2023 · 84e8ab1 · 84e8ab1
1 parent 538a2ea
commit 84e8ab1
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/batchflow/models/torch/base.py b/batchflow/models/torch/base.py
@@ -1551,7 +1551,10 @@ def aggregate_microbatches(self, outputs, chunked_outputs, chunk_sizes, single_o
                                          for chunk_output, chunk_size in zip(chunked_output, chunk_sizes)], dim=0)
                 result.append(output_)
             else:
-                result.append(np.mean([chunk_output.item() for chunk_output in chunked_output]))
+                if isinstance(chunked_output[0], np.ndarray):
+                    result.append(np.mean(chunked_output))
+                else:
+                    result.append(torch.mean(torch.stack(chunked_output)))
 
         if single_output:
             result = result[0]