Skip to content

GitLab

Commit 32ab028c authored by flyingdown's avatar flyingdown
Browse files

add --gpu-max-threads-per-block=1024 options

parent b10621d1
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
setup.py
View file @ 32ab028c
...@@ -169,6 +169,9 @@ if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4): ...@@ -169,6 +169,9 @@ if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4):
version_ge_1_5 = ["-DVERSION_GE_1_5"] version_ge_1_5 = ["-DVERSION_GE_1_5"]
version_dependent_macros = version_ge_1_1 + version_ge_1_3 + version_ge_1_5 version_dependent_macros = version_ge_1_1 + version_ge_1_3 + version_ge_1_5
if IS_ROCM_PYTORCH:
version_dependent_macros += ['--gpu-max-threads-per-block=1024']
if "--distributed_adam" in sys.argv or "--cuda_ext" in sys.argv: if "--distributed_adam" in sys.argv or "--cuda_ext" in sys.argv:
if "--distributed_adam" in sys.argv: if "--distributed_adam" in sys.argv:
sys.argv.remove("--distributed_adam") sys.argv.remove("--distributed_adam")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment