fix lufficc#206

lee-plus-plus · Jan 16, 2022 · a7f2937 · a7f2937 · qlhhahaha · Feb 26, 2022
1 parent 68dc0a2
commit a7f2937
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/ssd/utils/nms.py b/ssd/utils/nms.py
@@ -4,7 +4,10 @@
 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+')