Priyanka-Kumavat commited on
Commit
b09eafd
·
1 Parent(s): 3b14887

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
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