[fix]: fix wrapper comparison of pytorch version (open-mmlab#572)

wxzs5 · Nov 20, 2020 · 28260e5 · 28260e5
1 parent 202a57f
commit 28260e5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mmcv/cnn/bricks/wrappers.py b/mmcv/cnn/bricks/wrappers.py
@@ -30,7 +30,7 @@ def backward(ctx, grad):
 class Conv2d(nn.Conv2d):
 
     def forward(self, x):
-        if x.numel() == 0 and torch.__version__ <= '1.4':
+        if x.numel() == 0 and torch.__version__ <= '1.4.0':
             out_shape = [x.shape[0], self.out_channels]
             for i, k, p, s, d in zip(x.shape[-2:], self.kernel_size,
                                      self.padding, self.stride, self.dilation):
@@ -72,7 +72,7 @@ class MaxPool2d(nn.MaxPool2d):
 
     def forward(self, x):
         # PyTorch 1.6 does not support empty tensor inference yet
-        if x.numel() == 0 and torch.__version__ <= '1.6':
+        if x.numel() == 0 and torch.__version__ <= '1.6.0':
             out_shape = list(x.shape[:2])
             for i, k, p, s, d in zip(x.shape[-2:], _pair(self.kernel_size),
                                      _pair(self.padding), _pair(self.stride),