fix/typing(layer_set): Allow layer set read_proc to return single layer

ZettaAI · Aug 30, 2024 · 974ea6b · 974ea6b
1 parent d2834b2
commit 974ea6b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zetta_utils/layer/volumetric/layer_set/layer.py b/zetta_utils/layer/volumetric/layer_set/layer.py
@@ -35,7 +35,7 @@ class VolumetricLayerSet(
     read_procs: tuple[VolumetricSetDataProcT, ...] = ()
     write_procs: tuple[VolumetricSetDataWriteProcT, ...] = ()
 
-    def __getitem__(self, idx: UserVolumetricIndex) -> dict[str, npt.NDArray]:
+    def __getitem__(self, idx: UserVolumetricIndex) -> dict[str, npt.NDArray] | npt.NDArray:
         idx_backend = self.frontend.convert_idx(idx)
         return self.read_with_procs(idx=idx_backend)