How to add torch.compile() in LightningModule

[mactavish91]

I want to use pytorch2.0 to accelerate training, is there any method?

[akihironitta]

You can simply pass your LM to torch.compile(...):

compiled_model = torch.compile(your_lightning_module)
trainer.fit(compiled_model)

Issues/PRs/discussions labelled "pytorch 2.0" might have useful information for you :)