Update app.py

app.py

@@ -44,14 +44,16 @@ class ImageDescriptionDataset(Dataset):
 class LoRAModel(nn.Module):
     def __init__(self):
         super(LoRAModel, self).__init__()
-        self.backbone = models.resnet18(pretrained=True)  # Using a pre-trained ResNet18
+        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
 
     def forward(self, x):
-        x = self.backbone(x)
-        x = self.fc(x)
+        x = self.backbone(x)  # Pass through 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...")