fix lufficc#206

sieukim · Apr 14, 2022 · 2405763 · 2405763
1 parent 68dc0a2
commit 2405763
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/ssd/utils/nms.py b/ssd/utils/nms.py
@@ -4,7 +4,9 @@
 import torch
 import torchvision
 
-if torchvision.__version__ >= '0.3.0':
+from packaging import version
+
+if version.parse(torchvision.__version__) >= version.parse('0.3.0'):
     _nms = torchvision.ops.nms
 else:
     warnings.warn('No NMS is available. Please upgrade torchvision to 0.3.0+')