minor shape bug

py4dstem · Oct 8, 2024 · 175ec7d · 175ec7d
1 parent f37b976
commit 175ec7d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/py4DSTEM/process/phase/direct_ptychography.py b/py4DSTEM/process/phase/direct_ptychography.py
@@ -499,14 +499,15 @@ def preprocess(
             del _intensities
 
         self._intensities = copy_to_device(self._intensities, storage)
-        self._intensities_shape = np.array(self._intensities.shape[-2:])
+        self._intensities_shape = self._intensities.shape[-2:]
 
         # take FFT wrt real-space
         if vectorized_com_calculation:
             self._intensities_FFT = xp_storage.fft.fft2(self._intensities, axes=(0, 1))
         else:
-            self._intensities_FFT = xp_storage.empty_like(
-                self._intensities, dtype=xp_storage.complex64
+            self._intensities_FFT = xp_storage.empty(
+                self._intensities_shape + self._grid_scan_shape,
+                dtype=xp_storage.complex64,
             )
             # transpose to loop over cols faster
             intensities = self._intensities.transpose((2, 3, 0, 1))