Skip to content

Commit

Permalink
Fix torch.hub.list('ultralytics/yolov5') pathlib bug (ultralytics#3921
Browse files Browse the repository at this point in the history
)
  • Loading branch information
glenn-jocher committed Jul 7, 2021
1 parent 7dac0ab commit 3302d1f
Showing 1 changed file with 5 additions and 5 deletions.
10 changes: 5 additions & 5 deletions hubconf.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,9 @@
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
"""
from pathlib import Path

import torch

FILE = Path(__file__).absolute()


def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
"""Creates a specified YOLOv5 model
Expand All @@ -26,15 +23,18 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
Returns:
YOLOv5 pytorch model
"""
from pathlib import Path

from models.yolo import Model, attempt_load
from utils.general import check_requirements, set_logging
from utils.google_utils import attempt_download
from utils.torch_utils import select_device

check_requirements(requirements=FILE.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
file = Path(__file__).absolute()
check_requirements(requirements=file.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
set_logging(verbose=verbose)

save_dir = Path('') if str(name).endswith('.pt') else FILE.parent
save_dir = Path('') if str(name).endswith('.pt') else file.parent
path = (save_dir / name).with_suffix('.pt') # checkpoint path
try:
device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
Expand Down

0 comments on commit 3302d1f

Please sign in to comment.