Update model.py

Xtra-Computing · Jul 25, 2023 · 67c7c39 · 67c7c39
1 parent 66270c6
commit 67c7c39
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/model.py b/model.py
@@ -626,15 +626,15 @@ class ModelFedCon_noheader(nn.Module):
     def __init__(self, base_model, out_dim, n_classes, net_configs=None):
         super(ModelFedCon_noheader, self).__init__()
 
-        if base_model == "resnet50" or base_model == "resnet":
-            basemodel = models.resnet50(pretrained=False)
-            self.features = nn.Sequential(*list(basemodel.children())[:-1])
-            num_ftrs = basemodel.fc.in_features
-        elif base_model == "resnet18":
+        # if base_model == "resnet":
+        #     basemodel = models.resnet50(pretrained=False)
+        #     self.features = nn.Sequential(*list(basemodel.children())[:-1])
+        #     num_ftrs = basemodel.fc.in_features
+        if base_model == "resnet18":
             basemodel = models.resnet18(pretrained=False)
             self.features = nn.Sequential(*list(basemodel.children())[:-1])
             num_ftrs = basemodel.fc.in_features
-        elif base_model == "resnet50-cifar10" or base_model == "resnet50-cifar100" or base_model == "resnet50-smallkernel":
+        elif base_model == "resnet" or base_model == "resnet50-cifar10" or base_model == "resnet50-cifar100" or base_model == "resnet50-smallkernel":
             basemodel = ResNet50_cifar10()
             self.features = nn.Sequential(*list(basemodel.children())[:-1])
             num_ftrs = basemodel.fc.in_features