Fix: no_grad with AMP bug

[baskrahmer]

Fixes #20644

Note however that this would affect performance for other users, so the question is whether it is worth optimizing for this edge case that is fundamentally a torch bug.

cc @Borda

---

📚 Documentation preview 📚: https://pytorch-lightning--20921.org.readthedocs.build/en/20921/

[Borda]

Note however that this would affect performance for other users, so the question is whether it is worth optimizing for this edge case that is fundamentally a torch bug.

Then se shall report it and offer a fix in torch
Then, if it is accepted and released, we shall have a version switch in our codebase, so newer Torch versions won't need this compared to the old one. Does it...

BTW, have you measured the performance drop?
cc: @lantiga

[baskrahmer]

@Borda it is a long-standing issue in torch. I can try to make a fix if I have some time, but I think it could be complex.

But I agree with you that it should be fixed in torch ideally. Just wanted to open this PR to showcase what a workaround on our end would look like. Shall I close it?

I haven't measured the performance drop since it will vary strongly across architectures and probably also hardware setups.

[Borda]

@baskrahmer lets link the Torch issue also here for visibility 🐰

[baskrahmer]

@baskrahmer lets link the Torch issue also here for visibility 🐰

Sure, there's pytorch/pytorch#65766, pytorch/pytorch#112583 and pytorch/pytorch#105211