Honor gpu_id in process_slices

Author: trissim

src/arraybridge/slice_processing.py

@@ -7,9 +7,10 @@
 
 from arraybridge.converters import detect_memory_type
 from arraybridge.stack_utils import stack_slices, unstack_slices
+from arraybridge.utils import _get_device_id
 
 
-def process_slices(image, func, args, kwargs):
+def process_slices(image, func, args, kwargs, gpu_id=None):
     """
     Process a 3D array slice-by-slice using the provided function.
 
@@ -25,14 +26,18 @@ def process_slices(image, func, args, kwargs):
         func: Function to apply to each slice
         args: Positional arguments to pass to func
         kwargs: Keyword arguments to pass to func
+        gpu_id: Optional GPU device ID override. If not provided, attempts
+            to derive from the input image and falls back to 0.
 
     Returns:
         Processed 3D array, or tuple of (processed_3d_array, special_outputs...)
         if func returns tuples
     """
     # Detect memory type and use proper OpenHCS utilities
     memory_type = detect_memory_type(image)
-    gpu_id = 0  # Default GPU ID for slice processing
+    if gpu_id is None:
+        detected_gpu_id = _get_device_id(image, memory_type)
+        gpu_id = 0 if detected_gpu_id is None else detected_gpu_id
 
     # Unstack 3D array into 2D slices
     slices_2d = unstack_slices(image, memory_type, gpu_id)