From f41e3213642b856a578c603bbd0b03f3c5de1244 Mon Sep 17 00:00:00 2001
From: Jeffrey in Tiamat <thw20@ic.ac.uk>
Date: Sat, 6 Dec 2025 08:34:16 +0000
Subject: [PATCH] fix: lock optimum to 1.24.0

---
 setup.py                        | 2 +-
 src/chop/ir/graph/mase_graph.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index b37fb38ec..2de937aaf 100644
--- a/setup.py
+++ b/setup.py
@@ -61,7 +61,7 @@ def get_system():
     "kornia",
     "einops",
     "ghp-import",
-    "optimum",
+    "optimum==1.24.0",
     "pytest-profiling",
     "myst_parser",
     "pytest-cov",
diff --git a/src/chop/ir/graph/mase_graph.py b/src/chop/ir/graph/mase_graph.py
index 07de2cdc0..d2fd412ce 100644
--- a/src/chop/ir/graph/mase_graph.py
+++ b/src/chop/ir/graph/mase_graph.py
@@ -382,7 +382,7 @@ def from_checkpoint(
             MaseGraph: Loaded MaseGraph.
         """
         with open(f"{checkpoint}.pt", "rb") as f:
-            loaded_model = torch.load(f)
+            loaded_model = torch.load(f, weights_only=False)
 
         assert isinstance(
             loaded_model, fx.GraphModule