How to call torch.distributed.get_rank() in model building phase

[skjung0916]

I implemented pytorch lightning-based learning as follows.

---

dm = build_datamodule(config)
model = build_model(config)

trainer = Trainer(
...
accelerator="ddp",
...
)

trainer.fit(model, dm)

---

In this situation, in order to set different model parameters for each gpu process, distributed.get_rank() must be called at the stage of model building.
However, the trainer.fit function doesn't seem to be able to implement this because it requires an already built model.
I wonder if there is any other way to do this.

[ananthsub]

Hello, you are right that this currently isn't supported. I am working on adding this feature as part of this issue: #11922

Could you confirm that the issue and proposed solution meet your needs?