Debugged a torchvision version problem ()

Co-authored-by: wangshen <aibee@wangshen.local>
pull/609/head
Cathesilta 2021-12-03 10:30:52 +08:00 committed by GitHub
parent 9208930c11
commit fb1027de4c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 1 deletions
fastreid/modeling/backbones

View File

@ -4,7 +4,15 @@ from typing import Any, Callable, Dict, List, Optional, Sequence
import torch
from torch import nn, Tensor
from torch.nn import functional as F
from torchvision.models.utils import load_state_dict_from_url
#The style of importing Considers compatibility for the diversity of torchvision versions
try:
from torchvision.models.utils import load_state_dict_from_url
except ImportError:
try:
from torch.hub import load_state_dict_from_url
except ImportError:
from torch.utils.model_zoo import load_url as load_state_dict_from_url
from fastreid.layers import get_norm
from .build import BACKBONE_REGISTRY