Skip to content

Instantly share code, notes, and snippets.

@muellerzr

muellerzr/xla_script.py

Last active November 8, 2022 18:15
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save muellerzr/a85c9692101d47a9264a27fb5478225a to your computer and use it in GitHub Desktop.
Save muellerzr/a85c9692101d47a9264a27fb5478225a to your computer and use it in GitHub Desktop.
Download ZIP
Raw
xla_script.py
import torch_xla.core.xla_model as xm
import torch_xla.core.xla_env_vars as xenv
import os
def main(args=None):
print("I did something!")
if os.getenv(xenv.HOST_WORLD_SIZE, None) and xm.xrt_world_size() > 1:
xm.rendezvous("checking_out")
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment