diff --git a/ethosu/tensor_allocator/search_allocator.cpp b/ethosu/tensor_allocator/search_allocator.cpp
index c7c418a..5c7492b 100644
--- a/ethosu/tensor_allocator/search_allocator.cpp
+++ b/ethosu/tensor_allocator/search_allocator.cpp
@@ -68,7 +68,7 @@
             }
         }
     }
-    target_size = std::max(min_required_size, available_size);
+    target_size = std::max(min_required_size, size_limit);
     // Calculate the urgency of each live range
     lr_urgency.resize(lrs.size());
     for (size_t i = 0; i < lrs.size(); ++i) {
@@ -217,7 +217,8 @@
         // Pick any affecting live range.
         ix1 = turn_list[rng() % turn_list.size()];
     }
-    size_t ix2 = turn_list[rng() % turn_list.size() - 1];
+    // Note: turn_list has always at least 2 elements for bottlenecks
+    size_t ix2 = turn_list[rng() % (turn_list.size() - 1)];
     if (ix1 == ix2) {
         ix2 = turn_list[turn_list.size() - 1];
     }
diff --git a/ethosu/tensor_allocator/search_allocator.h b/ethosu/tensor_allocator/search_allocator.h
index 0ef8688..6c75015 100644
--- a/ethosu/tensor_allocator/search_allocator.h
+++ b/ethosu/tensor_allocator/search_allocator.h
@@ -98,10 +98,6 @@
      * The minimum possible size, assuming all live ranges can be perfectly allocated
      */
     uint32_t min_required_size;
-    /**
-     * The available size (input to algorithm).
-     */
-    uint32_t available_size;
     /** The algorithm stops once the target size has been achieved */
     uint32_t target_size;
     /** The highest end address of the best found allocation */
diff --git a/ethosu/tensor_allocator/test/test_tensor_allocator.py b/ethosu/tensor_allocator/test/test_tensor_allocator.py
index b9b547f..5350fde 100644
--- a/ethosu/tensor_allocator/test/test_tensor_allocator.py
+++ b/ethosu/tensor_allocator/test/test_tensor_allocator.py
@@ -16,15 +16,34 @@
 #
 # Description:
 # Unit tests for tensor_allocator.
+import pytest
+
 from ethosu import tensor_allocator
 
+test_data = [
+    ([(0, 100, 8000), (0, 1, 8016), (100, 110, 2000), (108, 110, 4000), (109, 110, 6000)], 16016),
+    (
+        [
+            (0, 23, 131072),
+            (4, 5, 65568),
+            (4, 9, 8192),
+            (8, 30, 15360),
+            (10, 11, 65568),
+            (10, 15, 4096),
+            (16, 17, 65552),
+            (16, 21, 2048),
+            (22, 23, 32784),
+            (22, 27, 1024),
+        ],
+        216096,
+    ),
+]
 
-def test_allocate():
+
+@pytest.mark.parametrize("lrs, expected_size", test_data)
+def test_allocate(lrs, expected_size):
     """Tests the search allocator"""
-    # Create an input that requires a search to produce a perfect allocation.
-    # Should result in size 16016
-    lrs = [(0, 100, 8000), (0, 1, 8016), (100, 110, 2000), (108, 110, 4000), (109, 110, 6000)]
     input = [x for lr in lrs for x in lr]
     res = tensor_allocator.allocate(input, 0)
     assert len(res) == len(lrs)
-    assert max(addr + lr[2] for addr, lr in zip(res, lrs)) == 16016
+    assert max(addr + lr[2] for addr, lr in zip(res, lrs)) == expected_size