use `pkg_resources` for PyTorch version comparison (#176)

use `pkg_resources` from `setuptools` to parse version strings, it is required by Pytorch >= 0.4.1 anyway
This commit is contained in:
Tan Jia Huei 2021-11-09 14:57:26 +08:00 committed by GitHub
parent 1a8b4b2899
commit 573315e83f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions

View File

@ -3,6 +3,7 @@ import os
import urllib
import warnings
from typing import Any, Union, List
from pkg_resources import packaging
import torch
from PIL import Image
@ -19,7 +20,7 @@ except ImportError:
BICUBIC = Image.BICUBIC
if [int(n.split("+")[0]) for n in torch.__version__.split(".")[:3]] < [1, 7, 1]:
if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
warnings.warn("PyTorch version 1.7.1 or higher is recommended")