Avoid a segmentation fault when clearing cached blocks (ROCm#297)

stanleytsang-amd · Oct 20, 2023 · 81c6746 · 81c6746
1 parent bc59a0d
commit 81c6746
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/hipcub/include/hipcub/backend/rocprim/util_allocator.hpp b/hipcub/include/hipcub/backend/rocprim/util_allocator.hpp
@@ -417,9 +417,7 @@ struct CachingDeviceAllocator
                     if (debug) _HipcubLog("\tDevice %d freed %lld bytes.\n\t\t  %lld available blocks cached (%lld bytes), %lld live blocks (%lld bytes) outstanding.\n",
                         device, (long long) block_itr->bytes, (long long) cached_blocks.size(), (long long) cached_bytes[device].free, (long long) live_blocks.size(), (long long) cached_bytes[device].live);
 
-                    cached_blocks.erase(block_itr);
-
-                    block_itr++;
+                    block_itr = cached_blocks.erase(block_itr);
                 }
 
                 // Unlock