From a7f2937c70f01b3bab2c7f48a1288afa9e377dc8 Mon Sep 17 00:00:00 2001
From: lee-plus-plus <88300311+lee-plus-plus@users.noreply.github.com>
Date: Sun, 16 Jan 2022 18:14:05 +0800
Subject: [PATCH] fix #206

---
 ssd/utils/nms.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/ssd/utils/nms.py b/ssd/utils/nms.py
index 5dd6d47b..894b1903 100644
--- 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+')