Update train_model.py

train_model.py

@@ -56,7 +56,8 @@ def train_model():
         print(f"Validation set size: {dataset_sizes['val']}")
         
         # Load a pre-trained model
-        model = models.resnet50(pretrained=True)
+        from torchvision.models import ResNet50_Weights
+        model = models.resnet50(weights=ResNet50_Weights.DEFAULT)
         
         # Modify the final layer for our number of classes
         num_ftrs = model.fc.in_features
@@ -75,7 +76,7 @@ def train_model():
         
         for epoch in range(num_epochs):
             print(f'Epoch {epoch+1}/{num_epochs}')
-            print('-' * 15)
+            print('-' * 10)
             
             # Each epoch has a training and validation phase
             for phase in ['train', 'val']: