diff --git a/cpp/src/loader/thread_batch_data_loader.cpp b/cpp/src/loader/thread_batch_data_loader.cpp
index 0b3ba3af4..9aff2d88c 100644
--- a/cpp/src/loader/thread_batch_data_loader.cpp
+++ b/cpp/src/loader/thread_batch_data_loader.cpp
@@ -210,7 +210,8 @@ uint8_t* ThreadBatchDataLoader::next_data()
 {
     if (num_workers_ == 0) // (location_len == 1 && batch_size == 1)
     {
-        // If it reads entire image with multi threads (using loader), release raster memory from batch data loader.
+        // If it reads entire image with multi threads (using loader), release raster memory from batch data loader
+        // by setting it to nullptr so that it will not be freed by ~ThreadBatchDataLoader (destructor).
         uint8_t* batch_raster_ptr = raster_data_[0];
         raster_data_[0] = nullptr;
         return batch_raster_ptr;
@@ -218,11 +219,7 @@ uint8_t* ThreadBatchDataLoader::next_data()
 
     if (processed_batch_count_ * batch_size_ >= location_len_)
     {
-        // Remove buffer items that are no longer needed.
-        for (size_t i = 0; i < buffer_item_len_; ++i)
-        {
-            raster_data_[i] = nullptr;
-        }
+        // If all batches are processed, return nullptr.
         return nullptr;
     }
 
diff --git a/python/pybind11/cucim_py.cpp b/python/pybind11/cucim_py.cpp
index bb8c37af8..62f041e2d 100644
--- a/python/pybind11/cucim_py.cpp
+++ b/python/pybind11/cucim_py.cpp
@@ -557,6 +557,10 @@ py::object py_associated_image(const CuImage& cuimg, const std::string& name, co
 py::object py_cuimage_iterator_next(CuImageIterator<CuImage>& it)
 {
     bool stop_iteration = (it.index() == it.size());
+    if (stop_iteration)
+    {
+        throw py::stop_iteration();
+    }
 
     // Get the next batch of images.
     ++it;
@@ -573,10 +577,6 @@ py::object py_cuimage_iterator_next(CuImageIterator<CuImage>& it)
         {
             _set_array_interface(cuimg_obj);
         }
-        if (stop_iteration)
-        {
-            throw py::stop_iteration();
-        }
         return cuimg_obj;
     }
 }