EMsurya commited on
Commit
69ef860
·
verified ·
1 Parent(s): 2e6c616

Update util/som.py

Browse files
Files changed (1) hide show
  1. util/som.py +1 -1
util/som.py CHANGED
@@ -3,7 +3,7 @@ from ultralytics import YOLO
3
  from PIL import Image
4
  import io
5
  import base64
6
- device = 'cuda'
7
 
8
  from PIL import Image, ImageDraw, ImageFont
9
  import numpy as np
 
3
  from PIL import Image
4
  import io
5
  import base64
6
+ device = "cuda" if torch.cuda.is_available() else "cpu"
7
 
8
  from PIL import Image, ImageDraw, ImageFont
9
  import numpy as np