AI-based_security_risk_identification_and_prevention_applications

Paused

qsitj commited on 12 days ago

Commit

e1fe61e

verified ·

1 Parent(s): 81d3e7f

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -31,6 +31,10 @@ detector50 = pipeline(model="facebook/detr-resnet-50")
 detector101 = pipeline(model="facebook/detr-resnet-101")
 COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",

 detector101 = pipeline(model="facebook/detr-resnet-101")
+if torch.cuda.is_available():
+    print("use cuda")
+    detector50.model.to('cuda')
+    detector101.model.to('cuda')
 COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",