Container -> Module

pytorch · Jan 17, 2017 · 13a0493 · 13a0493
1 parent 11821d6
commit 13a0493
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/torchvision/models/alexnet.py b/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(

diff --git a/torchvision/models/resnet.py b/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__()

diff --git a/torchvision/models/vgg.py b/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