Skip to content

Instantly share code, notes, and snippets.

@tchaton

tchaton/pytorch_lightning_xla_profiler.py

Created July 20, 2021 15:13
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 tchaton/d15ab3a5686ffb4e4dc7a6ed97b3d5a5 to your computer and use it in GitHub Desktop.
Save tchaton/d15ab3a5686ffb4e4dc7a6ed97b3d5a5 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
pytorch_lightning_xla_profiler.py
import pytorch_lightning as pl
model = MyLightningModule()
trainer = pl.Trainer(tpu_cores=8, profiler="xla")
trainer.fit(model)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment