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))