zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

py-torch: add M1 GPU support (#31283)

Browse Source
This commit is contained in:
Adam J. Stewart 2022-06-27 09:21:49 -07:00 committed by GitHub
parent 11d71ca85e
commit a6b0de3beb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
var/spack/repos/builtin/packages/py-torch/package.py
View File

@ -6,6 +6,7 @@
import os
import sys
from spack.operating_systems.mac_os import macos_version
from spack.package import *
@ -59,6 +60,7 @@ class PyTorch(PythonPackage, CudaPackage):
variant('kineto', default=True, description='Use Kineto profiling library', when='@1.8:')
variant('magma', default=not is_darwin, description='Use MAGMA', when='+cuda')
variant('metal', default=is_darwin, description='Use Metal for Caffe2 iOS build')
variant('mps', default=is_darwin and macos_version() >= Version('12.3'), description='Use MPS for macOS build', when='@1.12: platform=darwin')
variant('nccl', default=True, description='Use NCCL', when='+cuda platform=linux')
variant('nccl', default=True, description='Use NCCL', when='+cuda platform=cray')
variant('nccl', default=True, description='Use NCCL', when='+rocm platform=linux')
@ -373,6 +375,7 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
enable_or_disable('kineto')
enable_or_disable('magma')
enable_or_disable('metal')
enable_or_disable('mps')
enable_or_disable('breakpad')
enable_or_disable('nccl')