fix: bug

api.py

@@ -31,7 +31,7 @@ def load_annoy_index(
 import torch
 from torch import nn
 from torchvision import models, transforms
-from torchvision.models.vgg import VGG16_Weights
+# from torchvision.models.vgg import VGG16_Weights
 
 
 # Transform the image, so it becomes readable with the model
@@ -67,7 +67,8 @@ class FeatureExtractor(nn.Module):
     return out 
 
 # Initialize the model
-model = models.vgg16(weights=VGG16_Weights.DEFAULT)
+# model = models.vgg16(weights=VGG16_Weights.DEFAULT)
+model = models.vgg16(pretrained=True)
 new_model = FeatureExtractor(model)
 
 # Change the device to GPU