From 6632c0b6d8eb334b949facb07af4adfb5f75dd14 Mon Sep 17 00:00:00 2001
From: dzzz2001 <153698752+dzzz2001@users.noreply.github.com>
Date: Tue, 22 Oct 2024 17:49:33 +0800
Subject: [PATCH] fix nspin=4 error (#5319)

---
 source/module_hamilt_lcao/module_gint/gint_vl_gpu.cu | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/source/module_hamilt_lcao/module_gint/gint_vl_gpu.cu b/source/module_hamilt_lcao/module_gint/gint_vl_gpu.cu
index a406fd24fe..145f37417d 100644
--- a/source/module_hamilt_lcao/module_gint/gint_vl_gpu.cu
+++ b/source/module_hamilt_lcao/module_gint/gint_vl_gpu.cu
@@ -47,7 +47,8 @@ void gint_vl_gpu(hamilt::HContainer<double>* hRGint,
         checkCuda(cudaStreamCreate(&streams[i]));
     }
 
-    Cuda_Mem_Wrapper<double> grid_vlocal_g(hRGint->get_nnr(), 1, false);
+    const int nnrg = is_gamma_only ? hRGint->get_nnr() : gridt.nnrg;
+    Cuda_Mem_Wrapper<double> grid_vlocal_g(nnrg, 1, false);
     grid_vlocal_g.memset_device_sync();
 
     Cuda_Mem_Wrapper<double> dr_part(max_atom_per_z * 3, num_streams, true);
@@ -189,7 +190,7 @@ void gint_vl_gpu(hamilt::HContainer<double>* hRGint,
         checkCuda(cudaMemcpy(
             hRGint->get_wrapper(),
             grid_vlocal_g.get_device_pointer(),
-            hRGint->get_nnr() * sizeof(double),
+            nnrg * sizeof(double),
             cudaMemcpyDeviceToHost));
     }
     else
@@ -197,7 +198,7 @@ void gint_vl_gpu(hamilt::HContainer<double>* hRGint,
         checkCuda(cudaMemcpy(
             pvpR,
             grid_vlocal_g.get_device_pointer(),
-            hRGint->get_nnr() * sizeof(double),
+            nnrg * sizeof(double),
             cudaMemcpyDeviceToHost));
     }
     for (int i = 0; i < num_streams; i++)