What about an option where you assign topytorch_lightning.utilities.distributed, something likepytorch_lightning.utilities.distributed = pytorch_lightning.utilities.rank_zero, to make the code work without changes?

I want to confirm what you mean by this, I guess I'm a little confused. This snippet is from ddpm, do you mean something like this:

fromtorchvision.utilsimportmake_gridpytorch_lightning.utilities.distributed=pytorch_lightning.utilities.rank_zerofrompytorch_lightning.utilities.rank_zeroimportrank_zero_onlyfromomegaconfimportListConfig

If so, this doesn't work becausepytorch_lightning is not defined.

Another way I tried was this, but this doesn't seem to work:

importpytorch_lightning.utilities.rank_zeroasrank_zeroimportpytorch_lightning.utilities.distributedasdistributeddistributed=rank_zerodistributed.rank_zero_only()

However, this still does require some initial modification of thefrom pytorch_lightning.utilities.distributed import rank_zero_only line.
This, also doesn't seem to work as it still does try to initially import fromdistributed and results an error for me:

ModuleNotFoundError: No module named 'pytorch_lightning.utilities.distributed'Press any key to continue . . .

Perhaps, I am thinking about this or doing this incorrectly.

This is another way I thought of but doesn't work:

importpytorch_lightning.utilities.rank_zeroasrank_zerodefdistributed():returnrank_zerofrompytorch_lightning.utilities.distributedimportrank_zero_only

Something like this would work, but it does change the code (and is technically unnecessary since it doesn't do anything withdistributed:

importpytorch_lightning.utilities.rank_zeroasdistributedfrompytorch_lightning.utilities.rank_zeroimportrank_zero_only

If you prefer, I could instead replace this with a try catch/except like this which does work:

try:frompytorch_lightning.utilities.distributedimportrank_zero_onlyexcept:frompytorch_lightning.utilities.rank_zeroimportrank_zero_only

Sorry for the wall, this was just me trying to figure out a way of doing this

Something like this:

import pytorch_lightning.utilities.rank_zerosys.modules["pytorch_lightning.utilities.xxx"] = pytorch_lightning.utilities.rank_zerofrom pytorch_lightning.utilities.xxx import rank_zero_onlyrank_zero_only> <function lightning_utilities.core.rank_zero.rank_zero_only(fn: Callable) -> Callable>

That'd totally work! I just skipped over it initially because I wasn't sure if you wanted to allow importing sys or not because it's not already being imported, will adapt it to current edits

What I'm suggesting is to do this once, somewhere in code that's executed early (ie a function ininitialize_util called frominitialize.initialize()), and then all codebases will be able to usefrom pytorch_lightning.utilities.distributed import rank_zero_only.

You just gotta be careful and not replacepytorch_lightning.utilities.distributed if it exists already (which I assume is the case in the c uirrently used version of pytorch_lightning).

That'd make more sense! I'll revert and work on implementing something like that

Not sure if you want to keep the print statement in there or not as it's a bit unnecessary I feel like?

Sorry, I forgot to revert a few lines back, should be good now

Sorry, a bit of a mess, the import was actually definitely required but I moved it further down into the if