diff --git a/waveorder/cli/utils.py b/waveorder/cli/utils.py
index 67aad694..1f47109c 100644
--- a/waveorder/cli/utils.py
+++ b/waveorder/cli/utils.py
@@ -2,6 +2,7 @@
 from typing import Tuple
 
 import click
+import dask.array as da
 import numpy as np
 import torch
 from iohub.ngff import Position, open_ome_zarr
@@ -128,6 +129,7 @@ def apply_inverse_to_zyx_and_save(
             f"All values at t={t_idx} are zero or Nan, skipping reconstruction."
         )
         return
+    czyx_uint16_numpy = da.nan_to_num(czyx_uint16_numpy).compute()
 
     # convert to np.int32 (torch doesn't accept np.uint16), then convert to tensor float32
     czyx_data = torch.tensor(np.int32(czyx_uint16_numpy), dtype=torch.float32)