Merge pull request #674 from analysiscenter/release_fix

fixed torchmodel bug

Author: SergeyTsimfer

batchflow/__init__.py

@@ -27,4 +27,4 @@
 from .utils_telegram import TelegramMessage
 
 
-__version__ = '0.7.6'
+__version__ = '0.7.7'

batchflow/models/torch/base.py

@@ -1510,15 +1510,16 @@ def aggregate_microbatches(self, outputs, chunked_outputs, chunk_sizes, single_o
         result = []
         for i, _ in enumerate(outputs):
             # All tensors for current `output_name`
-            chunked_output = [chunk_outputs[i][:chunk_size]
-                              for chunk_outputs, chunk_size in zip(chunked_outputs, chunk_sizes)]
+            chunked_output = [chunk_outputs[i] for chunk_outputs in chunked_outputs]
             if chunked_output[0].size != 1:
                 if len(chunked_output) == 1:
                     output_ = chunked_output[0]
                 elif isinstance(chunked_output[0], np.ndarray):
-                    output_ = np.concatenate(chunked_output, axis=0)
+                    output_ = np.concatenate([chunk_output[:chunk_size]
+                                              for chunk_output, chunk_size in zip(chunked_output, chunk_sizes)], axis=0)
                 else:
-                    output_ = torch.cat(chunked_output, dim=0)
+                    output_ = torch.cat([chunk_output[:chunk_size]
+                                         for chunk_output, chunk_size in zip(chunked_output, chunk_sizes)], dim=0)
                 result.append(output_)
             else:
                 result.append(np.mean(chunked_output))