fix: memcpy_device_to_host size error

Author: OutisLi

source/lib/src/gpu/tabulate.cu

@@ -740,8 +740,8 @@ __global__ void tabulate_fusion_se_t_tebd_grad_fifth_order_polynomial(
   // gridDim.x -> nloc
 
   // Identify which atom and neighbor pair this thread is responsible for.
-  const int block_idx = blockIdx.x;
-  const int thread_idx = threadIdx.x;
+  const int_64 block_idx = blockIdx.x;
+  const int_64 thread_idx = threadIdx.x;
   const int ii = thread_idx / nnei_j;
   const int jj = thread_idx % nnei_j;
 
@@ -800,8 +800,8 @@ __global__ void tabulate_fusion_se_t_tebd_grad_grad_fifth_order_polynomial(
   // gridDim.x -> nloc
 
   // Identify which atom and neighbor pair this thread is responsible for.
-  const int block_idx = blockIdx.x;
-  const int thread_idx = threadIdx.x;
+  const int_64 block_idx = blockIdx.x;
+  const int_64 thread_idx = threadIdx.x;
   const int ii = thread_idx / nnei_j;
   const int jj = thread_idx % nnei_j;
 

source/lib/tests/test_tabulate_se_t_tebd.cc

@@ -718,8 +718,7 @@ TEST_F(TestTabulateSeTTebd, tabulate_fusion_se_t_tebd_gpu) {
       xyz_scatter_dev, table_dev, &table_info[0], em_x_dev, em_dev, nloc,
       nnei_i, nnei_j, last_layer_size);
 
-  deepmd::memcpy_device_to_host(&xyz_scatter[0], xyz_scatter_dev,
-                                xyz_scatter.size() * sizeof(double));
+  deepmd::memcpy_device_to_host(xyz_scatter_dev, xyz_scatter);
 
   deepmd::delete_device_memory(xyz_scatter_dev);
   deepmd::delete_device_memory(table_dev);
@@ -750,8 +749,7 @@ TEST_F(TestTabulateSeTTebd, tabulate_fusion_se_t_tebd_grad_gpu) {
       dy_dem_x_dev, table_dev, &table_info[0], em_x_dev, em_dev, dy_dev, nloc,
       nnei_i, nnei_j, last_layer_size);
 
-  deepmd::memcpy_device_to_host(&dy_dem_x[0], dy_dem_x_dev,
-                                dy_dem_x.size() * sizeof(double));
+  deepmd::memcpy_device_to_host(dy_dem_x_dev, dy_dem_x);
 
   deepmd::delete_device_memory(dy_dem_x_dev);
   deepmd::delete_device_memory(table_dev);