Model not saving when using distributed training

Hi -

I am trying to use the DECIModule with distributed training on 4 GPUs. Saving the model in the distributed case yields an empty module whereas when training on one GPU I save the anticipated model with sem_module.

Is there a known approach to saving the model properly in the distributed training case? I am using the pytorch default DDP strategy.

   trainer = pl.Trainer(
        accelerator="gpu",
        devices=4, # distribute training
        max_epochs=1000,
        fast_dev_run=test_run,
        callbacks=[TQDMProgressBar(refresh_rate=19), checkpoint_callback],
        enable_checkpointing=True,
    )

    # Training the model
    trainer.fit(lightning_module, datamodule=data_module)
    torch.save(lightning_module.sem_module, "model.pt")

Start with the DECIModule and the trainer setup shown in the issue, then compare the distributed and single-GPU paths around sem_module and checkpoint_callback. Reproduce the empty module result with the four-GPU DDP configuration and verify that the saved model contains the expected sem_module after training.