Support PyTorch-1.2.0

.travis.yml

@@ -8,6 +8,7 @@ python:
 
 # Supported PyTorch versions
 env:
+  - PYTORCH=1.2.0
   - PYTORCH=1.1.0
   - PYTORCH=1.0.1
   - PYTORCH=1.0.0

examples/amoebanet/__init__.py

@@ -1,6 +1,5 @@
 """An AmoebaNet-D implementation but using :class:`nn.Sequential`. :func:`amoebanetd`
 returns a :class:`nn.Sequential`.
-
 """
 from collections import OrderedDict
 from typing import Tuple
@@ -32,15 +31,16 @@ class ReLUConvBN(nn.Module):
     def __init__(self, in_planes: int, out_planes: int, kernel_size: int, stride: int,
                  padding: int, affine: bool = True):
         super().__init__()
-        self.op = nn.Sequential(
-            nn.ReLU(inplace=False),
-            nn.Conv2d(in_planes, out_planes, kernel_size,
-                      stride=stride, padding=padding, bias=False),
-            nn.BatchNorm2d(out_planes, affine=affine)
-        )
+        self.relu = nn.ReLU(inplace=False)
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size,
+                              stride=stride, padding=padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_planes, affine=affine)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
-        return self.op(x)
+        x = self.relu(x)
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
 
 
 class FactorizedReduce(nn.Module):
@@ -141,14 +141,12 @@ class Classifier(nn.Module):
 
     def __init__(self, channel_prev: int, num_classes: int):
         super().__init__()
-
         self.global_pooling = nn.AvgPool2d(7)
         self.classifier = nn.Linear(channel_prev, num_classes)
 
-    def forward(self, x: torch.Tensor) -> nn.Linear:  # type: ignore
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
         s1 = self.global_pooling(x[1])
         y = self.classifier(s1.view(s1.size(0), -1))
-
         return y
 
 

examples/resnet/bottleneck.py

@@ -1,6 +1,6 @@
 """A ResNet bottleneck implementation but using :class:`nn.Sequential`."""
 from collections import OrderedDict
-from typing import Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
 from torch import Tensor
 import torch.nn as nn
@@ -10,6 +10,11 @@
 Tensors = Tuple[Tensor, ...]
 TensorOrTensors = Union[Tensor, Tensors]
 
+if TYPE_CHECKING:
+    NamedModules = OrderedDict[str, nn.Module]
+else:
+    NamedModules = OrderedDict
+
 
 def conv3x3(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
     """3x3 convolution with padding"""
@@ -72,7 +77,7 @@ def bottleneck(inplanes: int,
                downsample: Optional[nn.Module] = None,
                ) -> nn.Sequential:
     """Creates a bottlenect block in ResNet as a :class:`nn.Sequential`."""
-    layers: Dict[str, nn.Module] = OrderedDict()
+    layers: NamedModules = OrderedDict()
     layers['twin'] = Twin()
 
     layers['conv1'] = Gutter(conv1x1(inplanes, planes))