This guide explains how to load the required Python and CUDA-accelerated PyTorch environment modules and safely submit deep learning jobs to the a40 partition.

To use PyTorch with GPU acceleration, you must load the specific Python package ecosystem tailored for the cluster's software ecosystem.

To see all available deep learning modules, run:

$ module avail unite/python/3.14/pytorch

Switch to a compute node environment under the high-performance a40 partition. PyTorch needs access to an NVIDIA A40 graphics card allocation to activate its CUDA backend:

$ srun --partition=a40 --gres=gpu:1 --cpus-per-task=4 --time=00:30:00 --pty bash

Once your shell prompt updates, you are safely running on an A40 GPU-enabled compute node.

Load the base Python 3.14 environment module together with the PyTorch deep learning framework module (which relies on the CUDA 12.8 backend wrappers):

$module load unite/python/3.14/python-3.14.0
$ module load unite/python/3.14/pytorch

To quickly check if PyTorch is loaded properly and can see the cluster's A40 GPUs, execute:

$ python3 -c "import torch; print('PyTorch Version:', torch.__version__); print('CUDA Available:', torch.cuda.is_available()); print('Device Name:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None')"

Create a basic neural processing script named `pytorch_test.py`. This example validates device selection, instantiates a simple neural network layer architecture, creates sample tensors, and verifies a basic forward propagation step on the GPU.

pytorch_test.py

import torch
import torch.nn as nn
 
# Explicitly detect and assign the cluster's GPU device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Target runtime device assigned: {device}")
 
if device.type == 'cuda':
    print(f"Active GPU Hardware: {torch.cuda.get_device_name(0)}")
else:
    print("WARNING: CUDA is not active. The script is falling back to CPU compute paths.")
 
# Define a minimal linear transformation model
model = nn.Sequential(
    nn.Linear(10, 50),
    nn.ReLU(),
    nn.Linear(50, 2)
).to(device)
 
print("Neural network layer architecture moved onto device memory.")
 
# Generate arbitrary mock tensor parameters directly on the device
inputs = torch.randn(32, 10).to(device)
print(f"Input batch tensor shape: {inputs.shape}")
 
# Execute forward processing pass
outputs = model(inputs)
print(f"Forward computation completed. Output matrix shape: {outputs.shape}")
 
# Simple validation calculation assertion
loss = outputs.sum()
print(f"Verification tracking metric: {loss.item():.4f}")
print("PyTorch model diagnostics ran successfully!")

You can safely test your pipeline code right here inside your interactive session environment by running:

$ python3 pytorch_test.py

Exit your interactive node and return to the primary login layer when your tests pass successfully:

$ exit

To launch deep learning tasks or train models headlessly in the background, create a batch submission file named `submit_pytorch.sh`:

submit_pytorch.sh

#!/bin/bash
#SBATCH --job-name=torch_dl_job
#SBATCH --partition=a40             # Target partition for the A40 nodes
#SBATCH --output=pytorch_job_%j.out
#SBATCH --error=pytorch_job_%j.err
#SBATCH --nodes=1                   # Request resources on a single node
#SBATCH --ntasks=1                  # Run a single task instance
#SBATCH --gres=gpu:1                # Request 1 A40 GPU allocation for tensor calculations
#SBATCH --cpus-per-task=4           # Request 4 CPU worker threads for background data loading
#SBATCH --time=01:00:00             # Max job runtime allocation (HH:MM:SS)
 
# Clean out inherited environment parameters
module purge
 
# Crucial: Load the exact same Python and PyTorch module pair
module load unite/python/3.14/python-3.14.0
module load unite/python/3.14/pytorch
 
# Run the PyTorch training or processing script
python3 pytorch_test.py

Submit the job to the scheduler queue:

$ sbatch submit_pytorch.sh

* Error: `ModuleNotFoundError: No module named 'torch'`
Solution: You missed loading the module framework or your active interactive session timed out. Run `module load unite/python/3.14/pytorch` first before calling your file. * Output says: `CUDA Available: False`
Solution: If the status check shows `False`, you either executed your test scripts on the shared login node, or forgot to define the GPU hardware resource parameter (`#SBATCH –gres=gpu:1`) inside your batch file script. Without these flags, Slurm will isolate your process to vanilla CPU cores.