fix: load torch weights with weights_only=True to prevent code execution (#728)

This commit is contained in:
Patrick Chilton authored and GitHub committed 2026-08-26 09:35:55 +08:00
1 parent 486ec538d3
commit 1087d4a6ec
1 file changed
+1 -1
@@ -62,6 +62,6 @@ class ModelLoader:
def _build_and_load_model(self, arch_config, model_path: Path):
model = BaseModel(arch_config)
state_dict = torch.load(model_path, map_location="cpu", weights_only=False)
state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
model.load_state_dict(state_dict)
return model