Enabling Unified Memory on Compute Cluster

dear all,

Some of my AF2 jobs with larger sequences are failing, for what I can only assume is memory issues. I have discovered it is possible to set up unified memory by adding the two flags below to my SLURM submission script, however the issue persists. I have been working with our IT department but we cannot resolve it. Is there any advice on what I might be doing wrong in my submission script, or other things I can check? thank you.

flags to add for unified memory:
```
export TF_FORCE_UNIFIED_MEMORY=1
export XLA_PYTHON_CLIENT_MEM_FRACTION="4.0"
```

my script (note I have removed actual paths from the script):
```
#!/bin/bash
#
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=8
#SBATCH --mem=200GB
#
#SBATCH --time=36:00:00
#SBATCH --no-requeue
#
#SBATCH --partition=gpu
#SBATCH --gres=gpu:4

MY_PROTEIN_PATH=(my_fasta_location.fasta)

echo $HOSTNAME

module load alphafold/2.2.4c

export OPENMM_CUDA_COMPILER=$(which nvcc)
export TF_FORCE_UNIFIED_MEMORY=1
export XLA_PYTHON_CLIENT_MEM_FRACTION="4.0"

python3 run_alphafold.py \
--model_preset=multimer \
--fasta_paths=$MY_PROTEIN_PATH \
--output_dir=$(dirname $MY_PROTEIN_PATH) \
--data_dir= \
--mgnify_database_path= \
--template_mmcif_dir= \
--max_template_date=2020-05-14 \
--obsolete_pdbs_path= \
--use_gpu_relax=true \
--bfd_database_path= \
--uniclust30_database_path= \
--uniref90_database_path= \
--pdb_seqres_database_path= \
--uniprot_database_path=

```