Spaces:
Runtime error
Runtime error
Satoru
commited on
Commit
•
f726e05
1
Parent(s):
20cfc19
fix: bug
Browse files
api.py
CHANGED
@@ -31,7 +31,7 @@ def load_annoy_index(
|
|
31 |
import torch
|
32 |
from torch import nn
|
33 |
from torchvision import models, transforms
|
34 |
-
from torchvision.models.vgg import VGG16_Weights
|
35 |
|
36 |
|
37 |
# Transform the image, so it becomes readable with the model
|
@@ -67,7 +67,8 @@ class FeatureExtractor(nn.Module):
|
|
67 |
return out
|
68 |
|
69 |
# Initialize the model
|
70 |
-
model = models.vgg16(weights=VGG16_Weights.DEFAULT)
|
|
|
71 |
new_model = FeatureExtractor(model)
|
72 |
|
73 |
# Change the device to GPU
|
|
|
31 |
import torch
|
32 |
from torch import nn
|
33 |
from torchvision import models, transforms
|
34 |
+
# from torchvision.models.vgg import VGG16_Weights
|
35 |
|
36 |
|
37 |
# Transform the image, so it becomes readable with the model
|
|
|
67 |
return out
|
68 |
|
69 |
# Initialize the model
|
70 |
+
# model = models.vgg16(weights=VGG16_Weights.DEFAULT)
|
71 |
+
model = models.vgg16(pretrained=True)
|
72 |
new_model = FeatureExtractor(model)
|
73 |
|
74 |
# Change the device to GPU
|