fix(ml): clear cache if bad zip file (#3600)

pull/4050/head
Mert 2023-08-08 13:33:13 +07:00 committed by GitHub
parent ddd4ec2d9e
commit fa57853bd2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions

@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
from pathlib import Path
from shutil import rmtree
from typing import Any
from zipfile import BadZipFile
from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf # type: ignore
@ -23,7 +24,7 @@ class InferenceModel(ABC):
loader = self.load if eager else self.download
try:
loader(**model_kwargs)
except (OSError, InvalidProtobuf):
except (OSError, InvalidProtobuf, BadZipFile):
self.clear_cache()
loader(**model_kwargs)