You can use SLURM to run your AI workload and the GPU to accelerate the workload.
Example on SLURM submission script for TensorFlow and Pytorch which request for 1 CPU task and 1 slice of A100:
#!/bin/env bash
#SBATCH --job-name=test_python
#SBATCH --time=10:00
#SBATCH --mem=16000
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:2g.20gb:1
#SBATCH --partition=gpu
module add cuda/cudnn/8.4.1 python/3.10.5
python train.py