Why is CPU Usage High in BatchRelaxer with mattersim-v1.0.0-5M?

[qchempku2017]

Dear community,

I've been running BatchRelaxer to relax 10,000 structures using the mattersim-v1.0.0-5M potential. However, I noticed that most of the computation happens on the CPU, while the GPU remains underutilized. The CPU seems to be the bottleneck (CPU: 32 core 64 threads, 64 GB RAM, GPU: V100 32G, n_atoms_per_batch=2048):

Since I'm not particularly familiar with mattersim's computing architecture, I’m curious about what might be causing this behavior. Given that gradient, energy, and Hessian computations primarily occur on the GPU, I would have expected the GPU, not the CPU, to be the main computational bottleneck.

Would it be possible to offload more work to the GPU to prevent the CPU from limiting the relaxation speed? If so, how could this be achieved?

Thanks in advance for any insights!