diff --git a/scripts/run_multimodal.py b/scripts/run_multimodal.py
index 231e340..1ed3555 100644
--- a/scripts/run_multimodal.py
+++ b/scripts/run_multimodal.py
@@ -113,7 +113,7 @@ def main(_):
   device = torch.device(_DEVICE.value)
   with _set_default_tensor_type(model_config.get_dtype()):
     model = gemma3_model.Gemma3ForMultimodalLM(model_config)
-    model.load_state_dict(torch.load(_CKPT.value)['model_state_dict'])
+    model.load_state_dict(torch.load(_CKPT.value, weights_only=True)['model_state_dict'])
     # model.load_weights(_CKPT.value)
     model = model.to(device).eval()
   print('Model loading done')