Disabling autocast for certain modules

LightningModule
1

Hi,
I was wondering what is the way in Lightning to disable mixed precision for certain sub-modules?
Is there a way to do this through callbacks?

Thanks

2

Hey

You can disable autocast over any section in your code through

with torch.autocast(enabled=False):
    # code here