diff options
| author | jario-jin <jariof@foxmail.com> | 2018-11-13 20:08:30 -0800 |
|---|---|---|
| committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2018-11-13 20:10:35 -0800 |
| commit | 0bedaf9cf6c2f393df3b505ba45c350b5e83fc86 (patch) | |
| tree | 730427d4ddf32a44d1af690532b50839fd9d8c2e /setup.py | |
| parent | 79ec5de3fc1b68befe9fb935aad354cf28f55db3 (diff) | |
| download | pytorch-0bedaf9cf6c2f393df3b505ba45c350b5e83fc86.tar.gz pytorch-0bedaf9cf6c2f393df3b505ba45c350b5e83fc86.tar.bz2 pytorch-0bedaf9cf6c2f393df3b505ba45c350b5e83fc86.zip | |
Update setup.py to support Nvidia TX2 (#13939)
Summary:
add platform.machine() == 'aarch64' for supporting Nvidia TX2
Pull Request resolved: https://github.com/pytorch/pytorch/pull/13939
Differential Revision: D13055834
Pulled By: soumith
fbshipit-source-id: 0fadc87adf9e6b796978ce743e824eb98b006856
Diffstat (limited to 'setup.py')
| -rw-r--r-- | setup.py | 5 |
1 files changed, 3 insertions, 2 deletions
@@ -193,10 +193,11 @@ IS_WINDOWS = (platform.system() == 'Windows') IS_DARWIN = (platform.system() == 'Darwin') IS_LINUX = (platform.system() == 'Linux') IS_PPC = (platform.machine() == 'ppc64le') +IS_ARM = (platform.machine() == 'aarch64') BUILD_PYTORCH = check_env_flag('BUILD_PYTORCH') -# ppc64le does not support MKLDNN -if IS_PPC: +# ppc64le and aarch64 do not support MKLDNN +if IS_PPC or IS_ARM: USE_MKLDNN = check_env_flag('USE_MKLDNN', 'OFF') else: USE_MKLDNN = check_env_flag('USE_MKLDNN', 'ON') |