More bugfixes in ML cupy/pycuda engine

ptycho · Feb 2, 2025 · dc2a127 · dc2a127
1 parent b54c8bb
commit dc2a127
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/ptypy/accelerate/cuda_cupy/engines/ML_cupy.py b/ptypy/accelerate/cuda_cupy/engines/ML_cupy.py
@@ -555,10 +555,11 @@ def new_grad(self):
             # local references
             ob = self.engine.ob.S[oID].data
             obg = ob_grad.S[oID].data
-            obf = ob_fln.S[oID].data
             pr = self.engine.pr.S[pID].data
             prg = pr_grad.S[pID].data
-            prf = pr_fln.S[pID].data
+            if self.engine.p.wavefield_precond:
+                obf = ob_fln.S[oID].data
+                prf = pr_fln.S[pID].data
 
             # Schedule w & I to device
             ev_w, w, data_w = self.engine.w_data.to_gpu(

diff --git a/ptypy/accelerate/cuda_pycuda/engines/ML_pycuda.py b/ptypy/accelerate/cuda_pycuda/engines/ML_pycuda.py
@@ -537,10 +537,11 @@ def new_grad(self):
             # local references
             ob = self.engine.ob.S[oID].data
             obg = ob_grad.S[oID].data
-            obf = ob_fln.S[oID].data
             pr = self.engine.pr.S[pID].data
             prg = pr_grad.S[pID].data
-            prf = pr_fln.S[pID].data
+            if self.engine.p.wavefield_precond:
+                obf = ob_fln.S[oID].data
+                prf = pr_fln.S[pID].data
 
             # Schedule w & I to device
             ev_w, w, data_w = self.engine.w_data.to_gpu(prep.weights, dID, qu_htod)