Fix incorrect argument type to torch.maximum.

ZettaAI · Oct 11, 2024 · 7b040e9 · 7b040e9
1 parent 095f02f
commit 7b040e9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py b/zetta_utils/mazepa_layer_processing/common/volumetric_apply_flow.py
@@ -95,7 +95,9 @@ def __call__(
                     intscn, subidx = src_idx.get_intersection_and_subindex(red_idx)
                     subidx_channels = [slice(0, res.shape[0])] + list(subidx)
                     with semaphore("read"):
-                        res[subidx_channels] = torch.maximum(res[subidx_channels], layer[intscn])
+                        res[subidx_channels] = torch.maximum(
+                            res[subidx_channels], convert.to_torch(layer[intscn], res.device)
+                        )
             else:
                 for src_idx, layer in zip(src_idxs, src_layers):
                     intscn, subidx = src_idx.get_intersection_and_subindex(red_idx)