SLURM batch example of python + conda + gpu usage.

gistfile1.txt

#!/bin/bash
#SBATCH --job-name=jobname        # Nombre del trabajo
#SBATCH --output=/home/user/logs/%j.log         # Nombre del output (%j se reemplaza por el ID del trabajo)
#SBATCH --error=/home/user/logs/%j.err          # Output de errores (opcional)
#SBATCH --ntasks=1                   # Correr 2 tareas
#SBATCH --cpus-per-task=2            # Numero de cores por tarea
#SBATCH --mem-per-cpu=0         # Memoria por proceso
#SBATCH --mail-type=END,FAIL         # Enviar eventos al mail (NONE, BEGIN, END, FAIL, ALL)
#SBATCH [email protected]    # El mail del usuario
#SBATCH --nodelist=node_name
#SBATCH --gres=gpu:4                 # Usar 2 GPUs (se pueden usar N GPUs de marca especifica de la manera --gres=gpu:marca:N)
#SBATCH --chdir=/home/user/code/

echo "Running"
eval "$(/home/user/miniconda3/bin/conda shell.bash hook)"
conda activate my_conda_env
which python
python -u /home/user/code/main.py --mode train --gpus [0,1,2,3]
echo "Finished with job $SLURM_JOBID"
date