easyblock = 'PythonPackage'

name = 'apex'
local_commit = '082f999'
version = '20210420'

homepage = 'https://github.com/nvidia/apex'
description = "A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch"

toolchain = {'name': 'fosscuda', 'version': '2020b'}

source_urls = ['https://github.com/NVIDIA/apex/archive/']
sources = [{
    'filename': SOURCE_TAR_GZ,
    'download_filename': '%s.tar.gz' % local_commit,
}]
checksums = ['a4a82b14ca95ec2068e85fdce5188d03c8eda92ac90a16bec46cd7e5e1534a53']

dependencies = [
    ('Python', '3.8.6'),
    ('PyTorch', '1.10.0'),
]

download_dep_fail = True
use_pip = True

installopts = "--no-cache-dir --global-option='--cpp_ext' --global-option='--cuda_ext' --global-option='--pyprof'"

sanity_pip_check = True

sanity_check_commands = ["python -c 'from apex import amp'"]

moduleclass = 'tools'