Update README.md

README.md

@@ -54,7 +54,6 @@ class CustomConvNeXtModel(nn.Module):
     def forward(self, x):
         return self.convnext(x)
 
-
 class CustomMobileNetModel(nn.Module):
     def __init__(self, weights=MobileNet_V2_Weights.DEFAULT, num_classes=2):
         super().__init__()
@@ -81,9 +80,9 @@ class CustomMobileNetModel(nn.Module):
         return self.mobilenet(x)
 
 class EnsembleModel(nn.Module):
-    def __init__(self, convnext_model, mobilenet_model, num_classes=2):
+    def __init__(self, resnet_model, mobilenet_model, num_classes=2):
         super().__init__()
-        self.convnext = convnext_model
+        self.resnet = resnet_model
         self.mobilenet = mobilenet_model
         self.fc = nn.Sequential(
             nn.Linear(num_classes * 2, 512), 
@@ -93,13 +92,12 @@ class EnsembleModel(nn.Module):
         )
 
     def forward(self, x):
-        convnext_out = self.convnext(x)
+        resnet_out = self.resnet(x)
         mobilenet_out = self.mobilenet(x)
-        combined = torch.cat((convnext_out, mobilenet_out), dim=1)
+        combined = torch.cat((resnet_out, mobilenet_out), dim=1)
         output = self.fc(combined)
         return output
 
-
 convnext_model = CustomConvNeXtModel()
 mobilenet_model = CustomMobileNetModel()
 ensemble_model = EnsembleModel(convnext_model, mobilenet_model)