Skip to content

Instantly share code, notes, and snippets.

@muellerzr

muellerzr/xla_scriptv2.py

Last active November 8, 2022 02:04
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/9a813285ee04e5007237dc6fe76053e3 to your computer and use it in GitHub Desktop.
Save muellerzr/9a813285ee04e5007237dc6fe76053e3 to your computer and use it in GitHub Desktop.
Download ZIP
Second version
Raw
xla_scriptv2.py
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.core.xla_model as xm
def _mp_fn(index):
print("I did something!")
xm.rendezvous('checking_out')
if __name__ == "__main__":
xmp.spawn(_mp_fn, args=(), nprocs=8)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment