Improve the gpu_device flag description

Based on the suggestion in https://github.com/google-deepmind/alphafold3/issues/493

PiperOrigin-RevId: 792089424
Change-Id: Iea8d0527725c2febbc5e09b619ca7ad1f778ddba
This commit is contained in:
Augustin Zidek
2025-08-07 03:52:57 -07:00
committed by Copybara-Service
parent 563896c75c
commit 2e2ffc10ab

View File

@ -228,9 +228,11 @@ _JAX_COMPILATION_CACHE_DIR = flags.DEFINE_string(
_GPU_DEVICE = flags.DEFINE_integer(
'gpu_device',
0,
'Optional override for the GPU device to use for inference. Defaults to the'
' 1st GPU on the system. Useful on multi-GPU systems to pin each run to a'
' specific GPU.',
'Optional override for the GPU device to use for inference, uses zero-based'
' indexing. Defaults to the 0th GPU on the system. Useful on multi-GPU'
' systems to pin each run to a specific GPU. Note that if GPUs are already'
' pre-filtered by the environment (e.g. by using CUDA_VISIBLE_DEVICES),'
' this flag refers to the GPU index after the filtering has been done.',
)
_BUCKETS = flags.DEFINE_list(
'buckets',