Update app.py

app.py

@@ -44,16 +44,18 @@ class ImageDescriptionDataset(Dataset):
 class LoRAModel(nn.Module):
     def __init__(self):
         super(LoRAModel, self).__init__()
-        self.backbone = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)  # Use the updated weights argument
-        # Adjust the final fully connected layer to match the output of the backbone
-        self.fc = nn.Linear(self.backbone.fc.in_features, 100)  # Placeholder output layer
+        self.backbone = models.resnet18(pretrained=True)  # Using a pre-trained ResNet18
+        
+        # Fixing the shape mismatch: Input size to the fc layer should match ResNet output
+        self.fc = nn.Linear(self.backbone.fc.in_features, 100)  # 100 is a placeholder for your output
+
+        # If you want to use LoRA, you will implement the low-rank adaptation mechanism here
 
     def forward(self, x):
-        x = self.backbone(x)  # Pass through the ResNet18 backbone
+        x = self.backbone(x)  # Extract features using the ResNet18 backbone
         x = self.fc(x)  # Apply the final fully connected layer
         return x
 
-
 # Function to train LoRA
 def train_lora(image_folder, metadata):
     print("Starting training process...")