update yolo.py

This commit is contained in:
Glenn Jocher committed 2020-06-18 19:33:02 -07:00
1 parent 245bd1d242
commit 80b82e8bde
1 file changed
+1 -1
+1 -1
View File
@@ -208,7 +208,7 @@ if __name__ == '__main__':
parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
opt = parser.parse_args()
opt.cfg = glob.glob('./**/' + opt.cfg, recursive=True)[0] # find file
opt.cfg = check_file(opt.cfg) # check file
device = torch_utils.select_device(opt.device)
# Create model