MultiTask Training on multi-gpus returns NaN and inf in model output during Validation phase
🐛 Bug
I'm trying to train a multihead ResNet on images. Train data size is 2700 and valid data size is 600. Each batch is [n, 3, 224, 224, 320] and normalized to have values [0,1]. I've already trained a single head resnet on this dataset many times and never encountered a problem with any datapoint so I am sure the data is good.
When I train the model using 2 heads on 1 GPU for 5 epochs, everything runs smoothly. If I try 2 GPUs, after the progress bar is done on train phase and starts validation phase, I start getting all NaNs from the feature extractor part of my network. I put in breakpoints to check the batch data and it has no NaNs or infs anywhere. I even made my validation data the same as my train data to check if the valid data was broken and feature extractor still outputs NaNs, even after just going through the same data during train phase with no error!
I can't find any pytorch lightning examples of MultiHead training, so if those exist please point me to those. I've followed the multi-gpu training docs for pytorch lightning to a T so have no idea what's going on. I'm also setting the random seed at beginning of program using pl.seed_everything(64). I'm using DDP and 16bit precision.

PyTorch Version (e.g., 1.0): 1.7.1+cu101
OS (e.g., Linux): Linux
How you installed PyTorch (conda, pip, source): pip
Build command you used (if compiling from source):
Python version: 3.7.6
CUDA/cuDNN version: 10.1
GPU models and configuration: 4 V100 on google cloud VM
Any other relevant information: 32 cores, 128GB mem

Additional context
Running ResNet101 on 3D images. Trying to train 3 classification heads (each head only looks at certain classes). Implemented my own MultiHeadLoss wrapper to calculate total loss from each head:
class MultiHeadLoss(nn.Module):
    """Sum of all losses."""

    def __init__(self, model_heads):
        super().__init__()
        self.model_heads = model_heads

    def forward(self, outputs_targets_dict):
        losses = []
        for head in self.model_heads:
            if head.name in outputs_targets_dict:
                output = outputs_targets_dict[head.name]["output"]
                target = outputs_targets_dict[head.name]["target"]
                head_loss = head.loss_fn(output, target)
                losses.append(head_loss)

        return sum(losses)

My network module looks like this:
class ModelTemplate(pl.LightningModule):
    def __init__(self, model_body=gin.REQUIRED, model_heads=gin.REQUIRED):
        super().__init__()
        self.model_body = model_body
        self.model_heads = nn.ModuleList()
        [self.model_heads.append(head(self.model_body.out_features)) for head in model_heads]

    def forward(self, x):
        features = self.model_body.forward(x)
        head_out = [head(features) for head in self.model_heads]
        return head_out

Each head looks like this:
class Classification(pl.LightningModule):
    def __init__(self, in_features, task_labels, name="classification"):
        super().__init__()
        self.register_buffer("task_labels", torch.Tensor(task_labels))
        self.name = name
        self.fc = nn.Linear(in_features, len(task_labels))
        self.loss_fn = WeightedFocalLoss(criterion=nn.CrossEntropyLoss(reduction="none"))
        self.activation_fn = nn.Softmax(dim=1)

    def forward(self, x):
        logits = self.fc(x)
        return logits

Run the train job like this:
trainer = Trainer(
        default_root_dir=results_dir,
        callbacks=callbacks,
        accelerator="ddp",
        gpus=gin.query_parameter("Trainer.gpus"),
        max_epochs=gin.query_parameter("Trainer.max_epochs"),
        precision=16,
        num_sanity_val_steps=False,
        logger=mlf_logger,
        profiler="simple",
    )
    trainer.fit(model=train_module, datamodule=datamodule)