Container -> Module

Author: soumith

torchvision/models/alexnet.py

@@ -10,7 +10,7 @@
 }
 
 
-class AlexNet(nn.Container):
+class AlexNet(nn.Module):
     def __init__(self, num_classes=1000):
         super(AlexNet, self).__init__()
         self.features = nn.Sequential(

torchvision/models/resnet.py

@@ -21,7 +21,7 @@ def conv3x3(in_planes, out_planes, stride=1):
                      padding=1, bias=False)
 
 
-class BasicBlock(nn.Container):
+class BasicBlock(nn.Module):
     expansion = 1
 
     def __init__(self, inplanes, planes, stride=1, downsample=None):
@@ -53,7 +53,7 @@ def forward(self, x):
         return out
 
 
-class Bottleneck(nn.Container):
+class Bottleneck(nn.Module):
     expansion = 4
 
     def __init__(self, inplanes, planes, stride=1, downsample=None):
@@ -92,7 +92,7 @@ def forward(self, x):
         return out
 
 
-class ResNet(nn.Container):
+class ResNet(nn.Module):
     def __init__(self, block, layers, num_classes=1000):
         self.inplanes = 64
         super(ResNet, self).__init__()

torchvision/models/vgg.py

@@ -7,7 +7,7 @@
 ]
 
 
-class VGG(nn.Container):
+class VGG(nn.Module):
     def __init__(self, features):
         super(VGG, self).__init__()
         self.features = features