Hugging Face's logo

Spaces:
qsitj
/
AI-based_security_risk_identification_and_prevention_applications
Paused

App Files Community
qsitj commited on
Commit
e1fe61e
·
verified ·
1 Parent(s): 81d3e7f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -0
app.py CHANGED
@@ -31,6 +31,10 @@ detector50 = pipeline(model="facebook/detr-resnet-50")
31
 
32
  detector101 = pipeline(model="facebook/detr-resnet-101")
33
 
 
 
 
 
34
 
35
 
36
  COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
 
31
 
32
  detector101 = pipeline(model="facebook/detr-resnet-101")
33
 
34
+ if torch.cuda.is_available():
35
+ print("use cuda")
36
+ detector50.model.to('cuda')
37
+ detector101.model.to('cuda')
38
 
39
 
40
  COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",