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