ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-02-07 21:35:00 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix torch CUDA 11.1 support

Browse Source
This commit is contained in:
allegroai 2020-11-26 01:14:36 +02:00
parent a8c543ef7b
commit 58eb5fbd5f
1 changed files with 25 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
trains_agent/helper/package/pytorch.py
View File

@ -82,6 +82,8 @@ class SimplePytorchRequirement(SimpleSubstitution):
92: 'https://download.pytorch.org/whl/cu92/torch_stable.html', 92: 'https://download.pytorch.org/whl/cu92/torch_stable.html',
100: 'https://download.pytorch.org/whl/cu100/torch_stable.html', 100: 'https://download.pytorch.org/whl/cu100/torch_stable.html',
101: 'https://download.pytorch.org/whl/cu101/torch_stable.html', 101: 'https://download.pytorch.org/whl/cu101/torch_stable.html',
102: 'https://download.pytorch.org/whl/cu102/torch_stable.html',
110: 'https://download.pytorch.org/whl/cu110/torch_stable.html',
} }
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
@ -117,18 +119,22 @@ class SimplePytorchRequirement(SimpleSubstitution):
@classmethod @classmethod
def get_torch_page(cls, cuda_version, nightly=False): def get_torch_page(cls, cuda_version, nightly=False):
# noinspection PyBroadException
try: try:
cuda = int(cuda_version) cuda = int(cuda_version)
except: except Exception:
cuda = 0 cuda = 0
if nightly: if nightly:
for c in range(cuda, max(-1, cuda-15), -1):
# then try the nightly builds, it might be there... # then try the nightly builds, it might be there...
torch_url = cls.nightly_page_lookup_template.format(cuda) torch_url = cls.nightly_page_lookup_template.format(c)
# noinspection PyBroadException
try: try:
if requests.get(torch_url, timeout=10).ok: if requests.get(torch_url, timeout=10).ok:
cls.torch_page_lookup[cuda] = torch_url print('Torch nightly CUDA {} download page found'.format(c))
return cls.torch_page_lookup[cuda], cuda cls.torch_page_lookup[c] = torch_url
return cls.torch_page_lookup[c], c
except Exception: except Exception:
pass pass
return return
@ -138,11 +144,14 @@ class SimplePytorchRequirement(SimpleSubstitution):
return cls.torch_page_lookup[cuda], cuda return cls.torch_page_lookup[cuda], cuda
# then try a new cuda version page # then try a new cuda version page
torch_url = cls.page_lookup_template.format(cuda) for c in range(cuda, max(-1, cuda-15), -1):
torch_url = cls.page_lookup_template.format(c)
# noinspection PyBroadException
try: try:
if requests.get(torch_url, timeout=10).ok: if requests.get(torch_url, timeout=10).ok:
cls.torch_page_lookup[cuda] = torch_url print('Torch CUDA {} download page found'.format(c))
return cls.torch_page_lookup[cuda], cuda cls.torch_page_lookup[c] = torch_url
return cls.torch_page_lookup[c], c
except Exception: except Exception:
pass pass