replace model with model.module

Author: ankitdhall

network/finetuner.py

@@ -137,34 +137,34 @@ def prepare_model(self, loading=False):
         # modify last layers based on the model being used
         if not loading:
             if self.model_name in ['alexnet', 'vgg']:
-                num_features = self.model.classifier[6].in_features
+                num_features = self.model.module.classifier[6].in_features
                 if isinstance(self.criterion, LastLevelCELoss):
-                    self.model.classifier[6] = nn.Linear(num_features, self.levels[-1])
+                    self.model.module.classifier[6] = nn.Linear(num_features, self.levels[-1])
                 elif isinstance(self.criterion, HierarchicalSoftmaxLoss):
-                    self.model.classifier[6] = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
+                    self.model.module.classifier[6] = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
                 else:
-                    self.model.classifier[6] = nn.Linear(num_features, self.n_classes)
+                    self.model.module.classifier[6] = nn.Linear(num_features, self.n_classes)
             elif 'resnet' in self.model_name:
-                num_features = self.model.fc.in_features
+                num_features = self.model.module.fc.in_features
                 if isinstance(self.criterion, LastLevelCELoss):
-                    self.model.fc = nn.Linear(num_features, self.levels[-1])
+                    self.model.module.fc = nn.Linear(num_features, self.levels[-1])
                 elif isinstance(self.criterion, HierarchicalSoftmaxLoss):
-                    self.model.fc = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
+                    self.model.module.fc = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
                 else:
-                    self.model.fc = nn.Linear(num_features, self.n_classes)
+                    self.model.module.fc = nn.Linear(num_features, self.n_classes)
         else:
             if self.model_name in ['alexnet', 'vgg']:
-                num_features = self.model.module.classifier[6].in_features
+                num_features = self.model.module.module.classifier[6].in_features
                 if isinstance(self.criterion, LastLevelCELoss):
-                    self.model.module.classifier[6] = nn.Linear(num_features, self.levels[-1])
+                    self.model.module.module.classifier[6] = nn.Linear(num_features, self.levels[-1])
                 elif isinstance(self.criterion, HierarchicalSoftmaxLoss):
-                    self.model.module.classifier[6] = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
+                    self.model.module.module.classifier[6] = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
                 else:
-                    self.model.module.classifier[6] = nn.Linear(num_features, self.n_classes)
+                    self.model.module.module.classifier[6] = nn.Linear(num_features, self.n_classes)
             elif 'resnet' in self.model_name:
-                num_features = self.model.module.fc.in_features
+                num_features = self.model.module.module.fc.in_features
                 if isinstance(self.criterion, LastLevelCELoss):
-                    self.model.module.fc = nn.Linear(num_features, self.levels[-1])
+                    self.model.module.module.fc = nn.Linear(num_features, self.levels[-1])
                 elif isinstance(self.criterion, HierarchicalSoftmaxLoss):
                     self.model.module.fc = HierarchicalSoftmax(labelmap=self.labelmap, input_size=num_features)
                 else: