diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index 5a84ca1..9c347e7 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -250,7 +250,8 @@ def get_images(idx):
         # once, but it avoids creating another self.patterns-sized array
         # as an intermediate step. This can be super important because
         # self.patterns can be more than half the available memory
-        nanomap_values = np.ones(self.translations.shape[0])
+        nanomap_values = t.ones(self.translations.shape[0],
+                                device=self.patterns.device)
 
         chunk_size = 10
         for i in range(0, self.translations.shape[0], chunk_size):
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 4d17f0b..d57eeaa 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -7,6 +7,7 @@
 import numpy as np
 import pytest
 import torch as t
+import matplotlib.pyplot as plt
 
 from cdtools.datasets import CDataset, Ptycho2DDataset
 from cdtools.tools import data as cdtdata
@@ -510,3 +511,32 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
         assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10, :])
 
         assert t.allclose(copied_dataset.translations, dataset.translations[10:-10, :])
+
+# This isn't actually slow, but it will fail by default if there is no
+# gpu on the machine
+@pytest.mark.slow
+def test_Ptycho2DDataset_inspect(ptycho_cxi_1, reconstruction_device, show_plot):
+    cxi, expected = ptycho_cxi_1
+    dataset = Ptycho2DDataset.from_cxi(cxi)
+
+    # Test for failure in several cases
+
+    # First, no additional plots
+    dataset.inspect(plot_mean_pattern=False, plot_mask=False)
+    plt.close('all')
+    
+    # Then, with additional plots
+    dataset.inspect(plot_mean_pattern=True, plot_mask=True)
+    plt.close('all')
+
+    # Finally, if data is on a special device like the GPU
+    dataset.to(device=reconstruction_device)
+    dataset.get_as(device=reconstruction_device)    
+    dataset.inspect(plot_mean_pattern=True, plot_mask=True)
+
+    if show_plot:
+        plt.show()
+    plt.close('all')
+
+
+