Initial commit (#5825)

deepmodeling · Jan 6, 2025 · 3d36512 · 3d36512
1 parent 4009f6a
commit 3d36512
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/source/module_basis/module_pw/module_fft/fft_bundle.cpp b/source/module_basis/module_pw/module_fft/fft_bundle.cpp
@@ -45,8 +45,15 @@ void FFT_Bundle::initfft(int nx_in,
 
     if (this->precision=="single")
     {
-        #ifndef __ENABLE_FLOAT_FFTW
-        float_define = false;
+        #if not defined (__ENABLE_FLOAT_FFTW)
+        if (this->device == "cpu"){
+            float_define = false;
+        }
+        #endif
+        #if defined(__CUDA) || defined (__ROCM)
+        if (this->device == "gpu"){
+            float_flag = float_define;
+        }
         #endif
         float_flag = float_define;
         double_flag = true;