Possible incorrect use of `and` with multiple context managers

[mariosgly]

Hi,

I noticed a small thing while looking through the training code and wanted to ask for your opinion. There are a few places where multiple context managers seem to be combined with and, for example:

with torch.cuda.amp.autocast() and torch.no_grad():

My understanding is that Python does not treat this as entering both context managers. Instead, it evaluates A() and B() first, so the with statement only receives one context manager, usually the second one if the first object is "true". So this may effectively behave like:

with torch.no_grad():

rather than:

with torch.cuda.amp.autocast(), torch.no_grad():

It does not seem like a huge issue, but it may mean that autocast() is not actually active in those blocks. In the opposite order, for example:

with torch.no_grad() and torch.cuda.amp.autocast():

it may mean that no_grad() is not active.

I may be missing something, so I wanted to check with you: do you think replacing these with comma-separated context managers would be the intended behavior?

Thanks!

[KrishnaParihar1]

Hey @mariosgly
I think this is a valid concern. In Python, with A() and B(): does not enter both context managers. The and expression is evaluated first and typically returns only the second context manager, since these objects are truthy. As a result, only one context manager is actually entered. If the intention is to enable both autocast() and no_grad(), the correct form would be:

with torch.cuda.amp.autocast(), torch.no_grad():

So it looks like autocast() or no_grad() may not be active in those blocks, depending on the order.

[Huaian666]

the more i dig into this the more i appreciate the design choices. clean codebase