Commit
·
b09eafd
1
Parent(s):
3b14887
Update app.py
Browse files
app.py
CHANGED
@@ -16,6 +16,8 @@ from tempfile import NamedTemporaryFile
|
|
16 |
# Create an FRCNN model instance with the same structure as the saved model
|
17 |
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(num_classes=91)
|
18 |
|
|
|
|
|
19 |
# Load the saved parameters into the model
|
20 |
model.load_state_dict(torch.load("frcnn_model.pth"))
|
21 |
|
|
|
16 |
# Create an FRCNN model instance with the same structure as the saved model
|
17 |
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(num_classes=91)
|
18 |
|
19 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
20 |
+
|
21 |
# Load the saved parameters into the model
|
22 |
model.load_state_dict(torch.load("frcnn_model.pth"))
|
23 |
|