glenn-jocher commited on
Commit
411842e
1 Parent(s): 87b094b

Fix `torch.hub.list('ultralytics/yolov5')` pathlib bug (#3921)

Browse files
Files changed (1) hide show
  1. hubconf.py +5 -5
hubconf.py CHANGED
@@ -4,12 +4,9 @@ Usage:
4
  import torch
5
  model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
6
  """
7
- from pathlib import Path
8
 
9
  import torch
10
 
11
- FILE = Path(__file__).absolute()
12
-
13
 
14
  def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
15
  """Creates a specified YOLOv5 model
@@ -26,15 +23,18 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
26
  Returns:
27
  YOLOv5 pytorch model
28
  """
 
 
29
  from models.yolo import Model, attempt_load
30
  from utils.general import check_requirements, set_logging
31
  from utils.google_utils import attempt_download
32
  from utils.torch_utils import select_device
33
 
34
- check_requirements(requirements=FILE.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
 
35
  set_logging(verbose=verbose)
36
 
37
- save_dir = Path('') if str(name).endswith('.pt') else FILE.parent
38
  path = (save_dir / name).with_suffix('.pt') # checkpoint path
39
  try:
40
  device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)
 
4
  import torch
5
  model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
6
  """
 
7
 
8
  import torch
9
 
 
 
10
 
11
  def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
12
  """Creates a specified YOLOv5 model
 
23
  Returns:
24
  YOLOv5 pytorch model
25
  """
26
+ from pathlib import Path
27
+
28
  from models.yolo import Model, attempt_load
29
  from utils.general import check_requirements, set_logging
30
  from utils.google_utils import attempt_download
31
  from utils.torch_utils import select_device
32
 
33
+ file = Path(__file__).absolute()
34
+ check_requirements(requirements=file.parent / 'requirements.txt', exclude=('tensorboard', 'thop', 'opencv-python'))
35
  set_logging(verbose=verbose)
36
 
37
+ save_dir = Path('') if str(name).endswith('.pt') else file.parent
38
  path = (save_dir / name).with_suffix('.pt') # checkpoint path
39
  try:
40
  device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device)