Multi-GPU Contrastive SSL: Runs fine with 1 or 2 GPUs, hangs with 4 GPUs

I'm trying to run the SimCLR contrastive learning demo as linked here (https://github.com/facebookresearch/SlowFast/blob/main/projects/contrastive_ssl/README.md) with the provided configs/contrastive_ssl/SimCLR_SlowR50_8x8.yaml file.

Using:
python tools/run_unet.py --cfg configs/contrastive_ssl/SimCLR_SlowR50_8x8.yaml DATA.PATH_TO_DATA_DIR path_to_dataset

If I set NUM_GPUS and BN.NUM_SYNC_DEVICES to 1 or 2, the training code runs fine and starts training. However, if I set these parameters to 4, even though I see the devices are properly allocated via nvidia-smi, the code seems to hang at inference on first input of the inputs to the model (e.g. stuck at the call to contrastive_forward(model, cfg, inputs, index, time, epoch_exact, scaler) in line 129 of train_net.py).

Any advice on how to proceed would be helpful, thanks