diff --git a/aom/aom_external_partition.h b/aom/aom_external_partition.h
index d944554..13a95f0 100644
--- a/aom/aom_external_partition.h
+++ b/aom/aom_external_partition.h
@@ -50,9 +50,16 @@
 
 /*!\brief Features pass to the external model to make partition decisions.
  * Specifically, features collected before NONE partition.
+ * Features "f" are used to determine:
+ * partition_none_allowed, partition_horz_allowed, partition_vert_allowed,
+ * do_rectangular_split, do_square_split
+ * Features "f_part2" are used to determine:
+ * prune_horz, prune_vert.
  */
 typedef struct aom_partition_features_before_none {
   float f[17]; /**< features to determine whether partition types allowed */
+  float
+      f_part2[25]; /**< features to determine whether partition types allowed */
 } aom_partition_features_before_none_t;
 
 /*!\brief Features pass to the external model to make partition decisions.
@@ -89,6 +96,18 @@
   float f[18]; /**< features to determine pruning 4-way partition */
 } aom_partition_features_ab_t;
 
+/*!\brief Feature id to tell the external model the current stage in partition
+ * pruning and what features to use to make decisions accordingly.
+ */
+typedef enum {
+  FEATURE_BEFORE_PART_NONE,
+  FEATURE_BEFORE_PART_NONE_PART2,
+  FEATURE_AFTER_PART_NONE,
+  FEATURE_AFTER_PART_SPLIT,
+  FEATURE_AFTER_PART_RECT,
+  FEATURE_AFTER_PART_AB
+} PART_FEATURE_ID;
+
 /*!\brief Features pass to the external model to make partition decisions.
  *
  * The encoder sends these features to the external model through
@@ -98,6 +117,7 @@
  * Once new features are finalized, bump the major version of libaom.
  */
 typedef struct aom_partition_features {
+  PART_FEATURE_ID id; /**< Feature ID to indicate active features */
   aom_partition_features_before_none_t
       before_part_none; /**< Features collected before NONE partition */
   aom_partition_features_none_t
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 351bc13..f4d0758 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -2022,6 +2022,7 @@
 
   // Setup features.
   aom_partition_features_t features;
+  features.id = FEATURE_BEFORE_PART_NONE;
   for (int i = 0; i < FEATURE_SIZE_SMS_SPLIT; ++i) {
     features.before_part_none.f[i] = features_from_motion[i];
   }
@@ -2058,8 +2059,9 @@
 
   // Setup features.
   aom_partition_features_t features;
+  features.id = FEATURE_BEFORE_PART_NONE_PART2;
   for (int i = 0; i < FEATURE_SIZE_SMS_PRUNE_PART; ++i) {
-    features.before_part_none.f[i] = features_from_motion[i];
+    features.before_part_none.f_part2[i] = features_from_motion[i];
   }
 
   // Send necessary features to the external model.
@@ -2094,6 +2096,7 @@
   if (!frame_is_intra_only(cm) && ext_part_controller->ready) {
     // Setup features.
     aom_partition_features_t features;
+    features.id = FEATURE_AFTER_PART_NONE;
     prepare_features_after_part_none(cpi, cm, x, sms_tree, part_search_state,
                                      pb_source_variance, &features);
 
@@ -2135,6 +2138,7 @@
   if (!frame_is_intra_only(cm) && cpi->ext_part_controller.ready) {
     // Setup features.
     aom_partition_features_t features;
+    features.id = FEATURE_AFTER_PART_SPLIT;
     prepare_features_after_split(cpi, x, sms_tree, part_search_state, best_rdc,
                                  partition_none_rdcost, partition_split_rdcost,
                                  split_block_rdcost, &features);
@@ -2189,6 +2193,7 @@
       ab_partition_allowed) {
     // Setup features.
     aom_partition_features_t features;
+    features.id = FEATURE_AFTER_PART_RECT;
     prepare_features_after_part_rect(cpi, x, pc_tree, bsize, best_rdcost,
                                      rect_part_rd, split_rd, &features);
 
@@ -2230,6 +2235,7 @@
   if (!frame_is_intra_only(cm) && ext_part_controller->ready) {
     // Setup features.
     aom_partition_features_t features;
+    features.id = FEATURE_AFTER_PART_AB;
     prepare_features_after_part_ab(cpi, x, bsize, part_ctx, best_rd,
                                    rect_part_rd, split_rd, pb_source_variance,
                                    mi_row, mi_col, &features);