small bug in dp_mask for gpu

py4dstem · May 31, 2024 · 6a48ab0 · 6a48ab0
1 parent c964ab7
commit 6a48ab0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/py4DSTEM/process/phase/phase_base_class.py b/py4DSTEM/process/phase/phase_base_class.py
@@ -686,6 +686,7 @@ def _calculate_intensities_center_of_mass(
 
                 # calculate CoM
                 if dp_mask is not None:
+                    dp_mask = copy_to_device(dp_mask, device)
                     intensities_mask = intensities * dp_mask
                 else:
                     intensities_mask = intensities