diff --git a/src/kbmod/run_search.py b/src/kbmod/run_search.py
index da582413..f0c05be5 100644
--- a/src/kbmod/run_search.py
+++ b/src/kbmod/run_search.py
@@ -131,13 +131,13 @@ def do_gpu_search(self, search, img_info, suggested_angle, post_process):
         if self.config["x_pixel_bounds"] and len(self.config["x_pixel_bounds"]) == 2:
             search.set_start_bounds_x(self.config["x_pixel_bounds"][0], self.config["x_pixel_bounds"][1])
         elif self.config["x_pixel_buffer"] and self.config["x_pixel_buffer"] > 0:
-            width = search.get_ImageStack().get_width()
+            width = search.get_imagestack().get_width()
             search.set_start_bounds_x(-self.config["x_pixel_buffer"], width + self.config["x_pixel_buffer"])
 
         if self.config["y_pixel_bounds"] and len(self.config["y_pixel_bounds"]) == 2:
             search.set_start_bounds_y(self.config["y_pixel_bounds"][0], self.config["y_pixel_bounds"][1])
         elif self.config["y_pixel_buffer"] and self.config["y_pixel_buffer"] > 0:
-            height = search.get_ImageStack().get_height()
+            height = search.get_imagestack().get_height()
             search.set_start_bounds_y(-self.config["y_pixel_buffer"], height + self.config["y_pixel_buffer"])
 
         # If we are using barycentric corrections, compute the parameters and
diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp
index 71dcdca2..ca273010 100644
--- a/src/kbmod/search/stack_search.cpp
+++ b/src/kbmod/search/stack_search.cpp
@@ -625,7 +625,9 @@ namespace search {
       .def("set_debug", &ks::set_debug, pydocs::DOC_StackSearch_set_debug)
       .def("filter_min_obs", &ks::filter_results, pydocs::DOC_StackSearch_filter_min_obs)
       .def("get_num_images", &ks::num_images, pydocs::DOC_StackSearch_get_num_images)
-      .def("get_imagestack", &ks::get_imagestack, pydocs::DOC_StackSearch_get_imagestack)
+      .def("get_imagestack", &ks::get_imagestack,
+          py::return_value_policy::reference_internal,
+          pydocs::DOC_StackSearch_get_imagestack)
       // Science Stamp Functions
       .def("get_stamps", &ks::get_stamps, pydocs::DOC_StackSearch_get_stamps)
       .def("get_median_stamp", &ks::get_median_stamp, pydocs::DOC_StackSearch_get_median_stamp)