Skip to content

Commit

Permalink
Replace distutils with packaging (#11)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed May 31, 2024
1 parent 5697e73 commit 5acf27a
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 4 deletions.
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ build-backend = "setuptools.build_meta"

[project]
name = "ultralytics-thop"
version = "0.0.2" # Placeholder version, needs to be dynamically set
version = "0.0.3" # Placeholder version, needs to be dynamically set
description = "A tool to count the FLOPs of PyTorch model."
readme = "README.md"
requires-python = ">=3.8"
Expand Down Expand Up @@ -57,6 +57,7 @@ classifiers = [
"Operating System :: Microsoft :: Windows",
]
dependencies = [
"packaging",
"torch",
]

Expand Down
6 changes: 3 additions & 3 deletions thop/profile.py
Original file line number Diff line number Diff line change
@@ -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 *
Expand All @@ -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__)
)
Expand Down Expand Up @@ -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})


Expand Down

0 comments on commit 5acf27a

Please sign in to comment.