_C = None def _lazy_import(): """ Make sure that CUDA versions match between the pytorch install and torchvision install """ global _C if _C is not None: return _C import torch from torchvision import _C as C _C = C if hasattr(_C, "CUDA_VERSION") and torch.version.cuda is not None: tv_version = str(_C.CUDA_VERSION) if int(tv_version) < 10000: tv_major = int(tv_version[0]) tv_minor = int(tv_version[2]) else: tv_major = int(tv_version[0:2]) tv_minor = int(tv_version[3]) t_version = torch.version.cuda t_version = t_version.split('.') t_major = int(t_version[0]) t_minor = int(t_version[1]) if t_major != tv_major or t_minor != tv_minor: raise RuntimeError("Detected that PyTorch and torchvision were compiled with different CUDA versions. " "PyTorch has CUDA Version={}.{} and torchvision has CUDA Version={}.{}. " "Please reinstall the torchvision that matches your PyTorch install." .format(t_major, t_minor, tv_major, tv_minor)) return _C