test_step hangs after one iteration when on multiple GPUs
🐛 Bug
When running the same code on a computer with 1 gpu, test_step runs as normal and logs what it should.
How ever on a node with 4 gpus, it hangs after 1 iteration!
Code sample
 images, masks = batch["image"], batch["mask"]
        if images.shape[1] != self.hparams.n_channels:
            raise AssertionError(
                f"Network has been defined with {self.n_channels} input channels, "
                f"but loaded images have {images.shape[1]} channels. Please check that "
                "the images are loaded correctly."
            )

        masks = (
            masks.type(torch.float32)
            if self.hparams.n_classes == 1
            else masks.type(torch.long)
        )

        masks_pred = self(images)  # Forward pass
        loss = self.loss_funciton(masks_pred, masks)
        result = pl.EvalResult(loss, checkpoint_on=loss)
        result.log("test_loss", loss, on_step=True, on_epoch=True, sync_dist=True)
        rand_idx = randint(0, self.hparams.batch_size - 1)
        onehot = torch.sigmoid(masks_pred[rand_idx]) > 0.5
        for tag, value in self.named_parameters():
            tag = tag.replace(".", "/")
            self.logger.experiment.add_histogram(tag, value, self.current_epoch)
        mask_grid = torchvision.utils.make_grid([masks[rand_idx], onehot], nrow=2)
        self.logger.experiment.add_image(
            "TEST - Target vs Predicted", mask_grid, self.current_epoch
        )
        alpha = 0.5
        image_grid = torchvision.utils.make_grid(
            [
                images[rand_idx],
                torch.clamp(
                    kornia.enhance.add_weighted(
                        src1=images[rand_idx],
                        alpha=1.0,
                        src2=onehot,
                        beta=alpha,
                        gamma=0.0,
                    ),
                    max=1.0,
                ),
            ]
        )
        self.logger.experiment.add_image(
            "TEST - Image vs Predicted", image_grid, self.current_epoch
        )
        pred = (torch.sigmoid(masks_pred) > 0.5).float()
        f1 = f1_score(pred, masks, self.hparams.n_classes + 1)
        rec = recall(pred, masks, self.hparams.n_classes + 1)
        pres = precision(pred, masks, self.hparams.n_classes + 1)
        result.log("test_f1", f1, on_epoch=True)
        result.log("test_recall", rec, on_epoch=True)
        result.log("test_precision", pres, on_epoch=True)

        return result
Expected behavior
I expect it to finish the testing-epoch.
Environment
Environment 1
CUDA:

GPU:
GeForce RTX 2070 SUPER
available: True
version: 10.2
Packages:
numpy: 1.19.2
pyTorch_debug: False
pyTorch_version: 1.6.0
pytorch-lightning: 0.9.0
tqdm: 4.49.0
System:
OS: Linux
architecture:
64bit
ELF
processor: x86_64
python: 3.6.9
version: #52~18.04.1-Ubuntu SMP Thu Sep 10 12:50:22 UTC 2020

Environment 2

CUDA:

GPU:

GeForce RTX 2080 Ti
GeForce RTX 2080 Ti
GeForce RTX 2080 Ti
GeForce RTX 2080 Ti


available:         True
version:           10.2


Packages:

numpy:             1.19.1
pyTorch_debug:     False
pyTorch_version:   1.6.0
pytorch-lightning: 0.9.0
tqdm:              4.49.0


System:

OS:                Linux
architecture:

64bit
ELF


processor:         x86_64
python:            3.8.0
version:           #208-Ubuntu SMP Sun Apr 5 23:45:10 UTC 2020