"CUBLAS_STATUS_NOT_INITIALIZED" error with fix pimd using dp model

[21al07se09t]

Hi everyone,
I'm trying to run a PIMD simulation with 32 beads in LAMMPS by fix pimd command, but I'm running into an Out of Memory (OOM) error.
My node has 2xNVIDIA GeForce RTX 3090 GPUs and 32 CPU cores.

I can perform simulations with 4 beads by following setup:

#!/bin/bash
#SBATCH -n 32
#SBATCH -N 1
#SBATCH -J 20-300
#SBATCH --gres=gpu:2
#SBATCH -p NVIDIAGeForceRTX3090
#SBATCH -e error.err
#SBATCH -o output.out

module purge
module add apps/lammps/2Aug2023-gpu-dp
export PATH=/public/software/apps/lammps/2Aug2023-dp/bin:$PATH

mpirun -np 32 lmp -partition 4x8 -in ./npt.in  -log ./npt.log

When I try to run a single PIMD simulation with 32 beads by following command, the job fails.
mpirun -np 32 lmp -partition 32x1 -in ./npt.in -log ./npt.log

I check error.err file and find that the OOM happens first, which then likely causes the cuBLAS initialization to fail:

2025-06-23 20:42:35.478254: E tensorflow/stream_executor/cuda/cuda_blas.cc:232] failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED
2025-06-23 20:42:35.478293: E tensorflow/stream_executor/cuda/cuda_blas.cc:234] Failure to initialize cublas may be due to OOM (cublas needs some free memory when you initialize it, and your deep-learning framework may have preallocated more than its fair share), or may be because this binary was not built with support for the GPU in your machine.
2025-06-23 20:42:35.478320: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at matmul_op_impl.h:451 : INTERNAL: Attempting to perform BLAS operation using StreamExecutor without BLAS support

So my question is: How can I resolve the problem I am encountering, and what is the recommended setup for running a PIMD simulation with a large number of beads (e.g., 32 or more)?

Thank you!

(full error.err file is attached)
error.txt