Multi-GPU doesn't work with SLURM

[ksikka]

From #138, multi-GPU doesn't work out of the box with SLURM

1. Add docs that you need to do the following

1. set --ntasks-per-node = 4 in your submit script, replacing 4 with the number of GPUs you want to use.

#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=4

2. use srun python ... instead of just python.

2. Fix code: LOCAL_RANK not set in SLURM

In addition to LOCAL_RANK, check SLURM_LOCALID in losses.py and dali.py. See Lightning-AI/pytorch-lightning#6797.