Skip to content

Commit

Permalink
[Fix] miss importing torchvision in torchscript wrapper (#1702)
Browse files Browse the repository at this point in the history
  • Loading branch information
grimoire authored and lvhan028 committed Mar 1, 2023
1 parent 60e7dc1 commit b4a977a
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion mmdeploy/backend/torchscript/wrapper.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
# Copyright (c) OpenMMLab. All rights reserved.
import importlib
import os.path as osp
from typing import Dict, Optional, Sequence, Union

import torch

from mmdeploy.utils import Backend
from mmdeploy.utils import Backend, get_root_logger
from mmdeploy.utils.timer import TimeCounter
from ..base import BACKEND_WRAPPER, BaseWrapper
from .init_plugins import get_ops_path
Expand Down Expand Up @@ -39,10 +40,20 @@ def __init__(self,
model: Union[str, torch.jit.RecursiveScriptModule],
input_names: Optional[Sequence[str]] = None,
output_names: Optional[Sequence[str]] = None):
logger = get_root_logger()

# load custom ops if exist
custom_ops_path = get_ops_path()
if osp.exists(custom_ops_path):
torch.ops.load_library(custom_ops_path)

# import torchvision for ops
try:
importlib.import_module('torchvision')
except Exception:
logger.warning(
'Can not import torchvision. '
'Models require ops in torchvision might not available.')
super().__init__(output_names)
self.ts_model = model
if isinstance(self.ts_model, str):
Expand Down

0 comments on commit b4a977a

Please sign in to comment.