TensorFlow

Guidance last updated: 20 October 2026

TensorFlow is an open source software library for numerical computation using data flow graphs. Nodes in the graph represent mathematical operations, while the graph edges represent the multidimensional data arrays (tensors) communicated between them. The flexible architecture allows you to deploy computation to one or more CPUs or GPUs in a desktop, server, or mobile device with a single API. TensorFlow was originally developed by researchers and engineers working on the Google Brain Team within Google’s Machine Intelligence research organization for the purposes of conducting machine learning and deep neural networks research, but the system is general enough to be applicable in a wide variety of other domains as well.

About TensorFlow on Stanage

Install TensorFlow with pip in your own Python virtual environment; see Python on Stanage. For Conda-managed environments, see Conda, then install TensorFlow using pip.

Note

GPU execution requires a GPU allocation; see Using GPUs on Stanage.

Depending on the TensorFlow wheel and GPU architecture, some GPU code may require just-in-time compilation from PTX during the first run. This can increase startup time; consult the TensorFlow hardware requirements.

Installing TensorFlow

First request a interactive session. For GPU installation and testing, request a GPU allocation.

The examples use Python 3.13, for which TensorFlow provides Linux wheels. Create a separate environment for your chosen installation:

module load Python/3.13.5-GCCcore-14.3.0
python -m venv "$HOME/.venvs/tensorflow-cpu"
source "$HOME/.venvs/tensorflow-cpu/bin/activate"
python -m pip install --upgrade pip
python -m pip install tensorflow-cpu

The and-cuda extra installs the NVIDIA user-space libraries required by TensorFlow. For this installation route, do not additionally load CUDA or cuDNN modules, which can introduce conflicting libraries. The cluster supplies the NVIDIA driver.

To select a particular TensorFlow release, specify its version, for example python -m pip install 'tensorflow[and-cuda]==2.21.0'. Consult the official installation guide for supported Python versions and GPU requirements.

In each new session or batch script, load the same Python module and activate the environment again before running your program.

After activating the GPU environment, expose its NVIDIA libraries and ptxas tool before starting TensorFlow. Repeat this setup in every GPU session and batch script:

NVIDIA_SITE_DIR=$(python -c 'import sysconfig; print(sysconfig.get_path("purelib")+"/nvidia")')
NVIDIA_LIB_DIRS=$(python -c 'import glob,os; print(":".join(sorted(glob.glob(os.environ["VIRTUAL_ENV"]+"/lib/python*/site-packages/nvidia/*/lib"))))')
export LD_LIBRARY_PATH="$NVIDIA_LIB_DIRS${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
export PATH="$NVIDIA_SITE_DIR/cuda_nvcc/bin:$PATH"

Testing your installation

Save this CPU check as tensorflow_check.py:

import tensorflow as tf

a = tf.constant([[1., 2., 3.], [4., 5., 6.]])
b = tf.constant([[1., 2.], [3., 4.], [5., 6.]])
c = tf.matmul(a, b)
tf.debugging.assert_near(c, [[22., 28.], [49., 64.]])
print(c)

For a GPU check, use the following instead. It requires a visible GPU and checks that the multiplication actually ran on it:

import tensorflow as tf

assert tf.config.list_physical_devices("GPU"), "No GPU detected"
tf.config.set_soft_device_placement(False)
with tf.device("/GPU:0"):
    a = tf.constant([[1., 2., 3.], [4., 5., 6.]])
    b = tf.constant([[1., 2.], [3., 4.], [5., 6.]])
    c = tf.matmul(a, b)
assert "GPU:0" in c.device, c.device
tf.debugging.assert_near(c, [[22., 28.], [49., 64.]])
print(c)

Batch usage

The following example batch submission script runs the GPU check above, saved as tensorflow_check.py, using the environment already created:

#!/bin/bash
#SBATCH --partition=gpu
#SBATCH --gres=gpu:a100:1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --mem-per-cpu=2000
#SBATCH --time=00:15:00

set -e
export SLURM_EXPORT_ENV=ALL
module load Python/3.13.5-GCCcore-14.3.0
source "$HOME/.venvs/tensorflow-gpu/bin/activate"
NVIDIA_SITE_DIR=$(python -c 'import sysconfig; print(sysconfig.get_path("purelib")+"/nvidia")')
NVIDIA_LIB_DIRS=$(python -c 'import glob,os; print(":".join(sorted(glob.glob(os.environ["VIRTUAL_ENV"]+"/lib/python*/site-packages/nvidia/*/lib"))))')
export LD_LIBRARY_PATH="$NVIDIA_LIB_DIRS${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
export PATH="$NVIDIA_SITE_DIR/cuda_nvcc/bin:$PATH"
srun python tensorflow_check.py

CUDA and cuDNN import errors

If TensorFlow cannot find its GPU libraries, check that the GPU environment is activated, tensorflow[and-cuda] was installed successfully, and the NVIDIA library/tool paths above are set. Use python -m pip check to check dependency consistency and follow the official GPU troubleshooting steps. Older installations that load CUDA/cuDNN modules must match the requirements of their TensorFlow release; see the tested build configurations.

For version-specific CUDA and cuDNN requirements, use the upstream tested build configurations linked above.

Training

The Research Software Engineering team has an introductory workshop on deep learning with the TensorFlow Keras framework.