Replace distutils with packaging

ultralytics · May 31, 2024 · 948de9d · 948de9d
1 parent 5697e73
commit 948de9d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -57,6 +57,7 @@ classifiers = [
     "Operating System :: Microsoft :: Windows",
 ]
 dependencies = [
+    "packaging",
     "torch",
 ]
 

diff --git a/thop/profile.py b/thop/profile.py
@@ -1,4 +1,4 @@
-from distutils.version import LooseVersion
+from packaging.version import Version
 
 from thop.rnn_hooks import *
 from thop.vision.basic_hooks import *
@@ -7,7 +7,7 @@
 # logger.setLevel(logging.INFO)
 from .utils import prGreen, prRed, prYellow
 
-if LooseVersion(torch.__version__) < LooseVersion("1.0.0"):
+if Version(torch.__version__) < Version("1.0.0"):
     logging.warning(
         "You are using an old version PyTorch {version}, which THOP does NOT support.".format(version=torch.__version__)
     )
@@ -61,7 +61,7 @@
     nn.PixelShuffle: zero_ops,
 }
 
-if LooseVersion(torch.__version__) >= LooseVersion("1.1.0"):
+if Version(torch.__version__) >= Version("1.1.0"):
     register_hooks.update({nn.SyncBatchNorm: count_normalization})