mirror of https://github.com/THUDM/CodeGeeX.git
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
24 lines
513 B
Python
24 lines
513 B
Python
import os
|
|
|
|
|
|
ENV_NAMES = ["CUDA_HOME", "LD_LIBRARY_PATH", "PATH", "TORCH_EXTENSIONS_DIR", "CUDA_LAUNCH_BLOCKING"]
|
|
|
|
|
|
def main():
|
|
s = ""
|
|
for name in ENV_NAMES:
|
|
if name in os.environ:
|
|
value = os.environ[name]
|
|
s += "{}={}\n".format(name, value)
|
|
print(f"{name}={value}")
|
|
else:
|
|
print(f"{name} is not set")
|
|
|
|
# write env vars to .deepspeed_env
|
|
with open(".deepspeed_env", "w") as f:
|
|
f.write(s)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|