fix nspin=4 error (#5319)

deepmodeling · Oct 22, 2024 · 6632c0b · 6632c0b
1 parent 3eac0e2
commit 6632c0b
Showing 1 changed file with 4 additions and 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
@@ -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,15 +190,15 @@ 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
     {
         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++)