(url, model_dir='./pretrained', map_location=None)
| 8 | |
| 9 | |
| 10 | def load_url(url, model_dir='./pretrained', map_location=None): |
| 11 | if not os.path.exists(model_dir): |
| 12 | os.makedirs(model_dir) |
| 13 | filename = url.split('/')[-1] |
| 14 | cached_file = os.path.join(model_dir, filename) |
| 15 | if not os.path.exists(cached_file): |
| 16 | sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file)) |
| 17 | urlretrieve(url, cached_file) |
| 18 | return torch.load(cached_file, map_location=map_location) |