Update README.md

README.md

@@ -29,3 +29,28 @@ state_dict = torch.load("resnet_gps_model.pth")
 resnet.load_state_dict(state_dict)
 resnet.eval()
 
+This is our customresnetmodel
+
+class CustomResNetModel(nn.Module):
+    def __init__(self, model_name="microsoft/resnet-18", num_classes=2):
+        super(CustomResNetModel, self).__init__()
+        # Load pre-trained ResNet from Hugging Face
+        self.resnet = AutoModelForImageClassification.from_pretrained(model_name)
+
+        # Adjust the classifier layer to output the desired number of classes
+        in_features = self.resnet.classifier[0].in_features  # Assuming the last layer is a Linear layer
+        self.resnet.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(in_features, num_classes)
+        )
+
+    def forward(self, x):
+        return self.resnet(x)
+
+    def save_model(self, save_path):
+        """Save model locally using the Hugging Face format."""
+        self.save_pretrained(save_path)
+
+    def push_model(self, repo_name):
+        """Push the model to the Hugging Face Hub."""
+        self.push_to_hub(repo_name)