Tensorflow¶

py-tensorflow / py-keras are not currently available in the 2026 stack

Neither py-tensorflow nor py-keras is built into the 2026 stack, so there is currently no shared module to load. Use your own virtual environment instead, as described below — the examples further down on this page assume you've already set one up.

Install your own TensorFlow with a virtual environment¶

You can install your own GPU-enabled TensorFlow directly on a login node using a plain Python virtual environment together with the gpu stack:

module load 2026 gpu python
python -m venv env
source env/bin/activate
pip install "tensorflow[and-cuda]"

The and-cuda extra pulls in matching nvidia-cudnn, nvidia-cublas, etc. pip packages, and also installs keras as a dependency (so both import tensorflow as tf and import keras as ks work — no separate py-keras needed).

Unlike PyTorch's wheels, TensorFlow does not automatically put these bundled CUDA libraries on its library search path — you need to point LD_LIBRARY_PATH at them yourself, every time you activate this environment, including inside your job scripts:

export LD_LIBRARY_PATH=$(python3 -c "
import os, nvidia
base = os.path.dirname(nvidia.__file__)
libdirs = [os.path.join(base, d, 'lib') for d in os.listdir(base) if os.path.isdir(os.path.join(base, d, 'lib'))]
print(':'.join(libdirs))
"):$LD_LIBRARY_PATH

Without this, TensorFlow silently falls back to CPU only (tf.config.list_physical_devices('GPU') returns [], with a generic "Cannot dlopen some GPU libraries" warning).

About the CUDA version

Do not module load cuda or cudnn alongside this environment — mixing an older module-provided cudnn with this pip-installed TensorFlow (which expects cuDNN ≥ 9.3) causes CUDNN_STATUS_INTERNAL_ERROR. The and-cuda extra is fully self-contained once LD_LIBRARY_PATH is set as above.

Loading the gpu stack before creating the environment is what makes this work from a login node without a GPU card present, and the resulting env will run unmodified on the GPU compute nodes (as long as LD_LIBRARY_PATH is set there too). See Installing GPU-enabled packages for more details on this workflow.

Checking GPU availability¶

With the env virtual environment set up as above, let's try to run the following minimal example of a Tensorflow job (howmanygpus.py):

howmanygpus.py

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

This can be submitted to the queue with the following submission script, which activates env and sets LD_LIBRARY_PATH before running:

howmanygpus.slurm

#!/bin/bash -l
#
#SBATCH --job-name="tensorflow/howmanygpus"
#SBATCH --output=howmanygpus.out
#SBATCH --time=00:02:00
#SBATCH --partition=gpu-a100-small
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --gpus-per-task=2
#SBATCH --mem-per-cpu=1G
# make sure to add your account!
##SBATCH --account=<what>-<faculty>-<group>

# requires a virtual environment "env" already set up as described in
# "Install your own TensorFlow with a virtual environment" in the docs
module load 2026 gpu python
source env/bin/activate

export LD_LIBRARY_PATH=$(python3 -c "
import os, nvidia
base = os.path.dirname(nvidia.__file__)
libdirs = [os.path.join(base, d, 'lib') for d in os.listdir(base) if os.path.isdir(os.path.join(base, d, 'lib'))]
print(':'.join(libdirs))
"):$LD_LIBRARY_PATH

srun python howmanygpus.py

Keras classification example¶

A slightly more advanced example that actually "does something" is given below. In the job script, we clone a benchmark data set onto the local SSD of a GPU node, and then perform a classification training (example taken from the TensorFlow Keras tutorial:

classifition.py

# TensorFlow and tf.keras
import tensorflow as tf
import keras as ks

# Find GPUs
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

# Helper libraries
import numpy as np

print(tf.__version__)

# Import the Fashion MNIST dataset
# 60,000 images are used to train the network and 10,000 images to evaluate:
fashion_mnist = ks.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

print('Train images shape is: '+ str(train_images.shape))
print('Length of train labels is: '+ str(len(train_labels)))
print('Test images shape is: '+ str(test_images.shape))
print('Length of test labels is: '+ str(len(test_labels)))

# Scale these values to a range of 0 to 1 before feeding them to the neural network model.
train_images = train_images / 255.0
test_images = test_images / 255.0


# Set up the layers
model = ks.Sequential([
    ks.layers.Input(shape=(28,28)),
    ks.layers.Flatten(),
    ks.layers.Dense(128, activation='relu'),
    ks.layers.Dense(10)
])

# Print Summary
model.summary()

# Compile the model
model.compile(optimizer='adam',
              loss=ks.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# Train the model
model.fit(train_images, train_labels, epochs=10)

# Evaluate accuracy
test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)
print('\nTest accuracy:', test_acc)

And here is the corresponding job script:

classifition.slurm

#!/bin/bash -l
#
#SBATCH --job-name="tensorflow/classification"
#SBATCH --output=classification.out
#SBATCH --partition=gpu-a100-small
#SBATCH --time=00:10:00
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --gpus-per-task=1
#SBATCH --mem-per-cpu=3G
# make sure to add your account!
##SBATCH --account=<what>-<faculty>-<group>

# requires a virtual environment "env" already set up as described in
# "Install your own TensorFlow with a virtual environment" in the docs
module load 2026 gpu python
source env/bin/activate

export LD_LIBRARY_PATH=$(python3 -c "
import os, nvidia
base = os.path.dirname(nvidia.__file__)
libdirs = [os.path.join(base, d, 'lib') for d in os.listdir(base) if os.path.isdir(os.path.join(base, d, 'lib'))]
print(':'.join(libdirs))
"):$LD_LIBRARY_PATH

#note: before submitting this script, you need to clone the fashion mnist
# repository, which should be done on the login node. Compute/GPU nodes do not
# have an internet connection.
## git clone https://github.com/zalandoresearch/fashion-mnist fashion_mnist

# for fastest possible I/O on a single node, we will work
# on the node-local SSD
RUNDIR=/tmp/${SLURM_JOBID}
mkdir ${RUNDIR}
cp classification.py ${RUNDIR}
rsync -a fashion_mnist ${RUNDIR}
cd ${RUNDIR}

srun python classification.py

# always clean up after yourself...
rm -rf ${RUNDIR}