diff --git a/miplearn/io.py b/miplearn/io.py
index 635be8c..ccad431 100644
--- a/miplearn/io.py
+++ b/miplearn/io.py
@@ -86,7 +86,8 @@ def read_pkl_gz(filename: str) -> Any:
 
 def _to_h5_filename(data_filename: str) -> str:
     output = f"{data_filename}.h5"
-    output = output.replace(".pkl.gz.h5", ".h5")
+    output = output.replace(".gz.h5", ".h5")
+    output = output.replace(".json.h5", ".h5")
     output = output.replace(".pkl.h5", ".h5")
     output = output.replace(".jld2.h5", ".h5")
     return output