mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-03 04:40:06 +06:00
8fff07ded0
* [run-slow] cohere * [run-slow] cohere * [run-slow] cohere --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
27 lines
861 B
Python
27 lines
861 B
Python
"""A simple script to set flexibly CUDA_VISIBLE_DEVICES in GitHub Actions CI workflow files."""
|
|
|
|
import argparse
|
|
import os
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--test_folder",
|
|
type=str,
|
|
default=None,
|
|
help="The test folder name of the model being tested. For example, `models/cohere`.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# `test_eager_matches_sdpa_generate` for `cohere` needs a lot of GPU memory!
|
|
# This depends on the runners. At this moment we are targeting our AWS CI runners.
|
|
if args.test_folder == "models/cohere":
|
|
cuda_visible_devices = "0,1,2,3"
|
|
elif "CUDA_VISIBLE_DEVICES" in os.environ:
|
|
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
|
|
else:
|
|
cuda_visible_devices = "0"
|
|
|
|
print(cuda_visible_devices)
|