lijunyu commited on
Commit
0dec8ff
1 Parent(s): 4dd75b3

path fix for torch hub

Browse files
Files changed (1) hide show
  1. hubconf.py +5 -1
hubconf.py CHANGED
@@ -6,6 +6,9 @@ Usage:
6
  """
7
 
8
  dependencies = ['torch', 'yaml']
 
 
 
9
  import torch
10
 
11
  from models.yolo import Model
@@ -24,7 +27,8 @@ def create(name, pretrained, channels, classes):
24
  Returns:
25
  pytorch model
26
  """
27
- model = Model('models/%s.yaml' % name, channels, classes)
 
28
  if pretrained:
29
  ckpt = '%s.pt' % name # checkpoint filename
30
  google_utils.attempt_download(ckpt) # download if not found locally
 
6
  """
7
 
8
  dependencies = ['torch', 'yaml']
9
+
10
+ import os
11
+
12
  import torch
13
 
14
  from models.yolo import Model
 
27
  Returns:
28
  pytorch model
29
  """
30
+ config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name)
31
+ model = Model(config, channels, classes)
32
  if pretrained:
33
  ckpt = '%s.pt' % name # checkpoint filename
34
  google_utils.attempt_download(ckpt) # download if not found locally