Update app.py

app.py

@@ -2,7 +2,8 @@ import os
 import zipfile
 import numpy as np
 import torch
-from transformers import ViTForImageClassification, AdamW
+from transformers import ViTForImageClassification
+from torch.optim import AdamW
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import streamlit as st
@@ -80,11 +81,12 @@ class CustomImageDataset(Dataset):
         label = self.labels[idx]
         return image, label
 
+# Training function
 # Training function
 def fine_tune_model(train_loader):
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=3)
     model.train()
-    optimizer = AdamW(model.parameters(), lr=1e-4)
+    optimizer = AdamW(model.parameters(), lr=1e-4)  # Use PyTorch's AdamW
     criterion = torch.nn.CrossEntropyLoss()
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)