mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 21:10:55 +08:00
fix: load torch weights with weights_only=True to prevent code execution (#728)
This commit is contained in:
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
|
||||
Reference in new issue
Block a user