Merge pull request #87 from neuro-ml/develop

Compatibility for custom PatchCombiner
neuro-ml · Jan 22, 2024 · 62eb807 · 62eb807
2 parents 778d36b + 5580873
commit 62eb807
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/dpipe/__version__.py b/dpipe/__version__.py
@@ -1 +1 @@
-__version__ = '0.3.2'
+__version__ = '0.3.3'
diff --git a/dpipe/predict/shape.py b/dpipe/predict/shape.py
@@ -123,9 +123,13 @@ def wrapper(x, *args, **kwargs):
             else:
                 patches = pmap(predict, input_patches, *args, **kwargs)
 
+            prediction_kwargs = {'combiner': combiner, 'get_boxes': get_boxes}
+            if not use_torch:
+                prediction_kwargs['use_torch'] = use_torch
+
             prediction = combine(
                 patches_wrapper(patches), extract(x.shape, input_axis), local_stride, axis,
-                combiner=combiner, get_boxes=get_boxes, use_torch=use_torch
+                **prediction_kwargs
             )
 
             if valid: