diff --git a/bayesflow/adapters/transforms/nan_to_num.py b/bayesflow/adapters/transforms/nan_to_num.py
index 19a94e5ae..9c047da0c 100644
--- a/bayesflow/adapters/transforms/nan_to_num.py
+++ b/bayesflow/adapters/transforms/nan_to_num.py
@@ -80,6 +80,8 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         data = data.copy()
 
         # Retrieve mask and values to reconstruct NaNs
+        if self.key not in data.keys():
+            return data
         values = data[self.key]
 
         if not self.return_mask: