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:
parent
1a8b4b2899
commit
573315e83f
|
@ -3,6 +3,7 @@ import os
|
||||||
import urllib
|
import urllib
|
||||||
import warnings
|
import warnings
|
||||||
from typing import Any, Union, List
|
from typing import Any, Union, List
|
||||||
|
from pkg_resources import packaging
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
@ -19,7 +20,7 @@ except ImportError:
|
||||||
BICUBIC = Image.BICUBIC
|
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")
|
warnings.warn("PyTorch version 1.7.1 or higher is recommended")
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue