diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py
index 67ae536aa..a66e91911 100644
--- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py
+++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py
@@ -19,6 +19,7 @@
 """
 
 import argparse
+import ast
 import os
 import time
 from contextlib import contextmanager
@@ -82,8 +83,11 @@ def create_state_spec_from_checkpoint(ckpt_path: str):
             continue
 
         if isinstance(value, dict) and "shape" in value and "dtype" in value:
-            # pylint: disable=eval-used
-            shape = eval(value["shape"]) if isinstance(value["shape"], str) else value["shape"]
+            shape = (
+                ast.literal_eval(value["shape"])
+                if isinstance(value["shape"], str)
+                else value["shape"]
+            )
             dtype_str = value["dtype"]
 
             # Convert dtype string to jax dtype