1 files changed, 32 insertions, 15 deletions

diff --git a/src/graph/mutators/NodeFusionMutator.cpp b/src/graph/mutators/NodeFusionMutator.cpp
index 2e893c2e07..39209d2f49 100644
--- a/src/graph/mutators/NodeFusionMutator.cpp
+++ b/src/graph/mutators/NodeFusionMutator.cpp
@@ -54,28 +54,45 @@ void fuse_batch_norm_with_activation(Graph &g)
                 auto *bn_node  = arm_compute::utils::cast::polymorphic_downcast<BatchNormalizationLayerNode *>(output_edge->producer());
                 auto *act_node = arm_compute::utils::cast::polymorphic_downcast<ActivationLayerNode *>(output_edge->consumer());
 
-                // Get driving nodes of activation node
-                std::vector<NodeIdxPair> act_driving_nodes;
-                for(auto &act_output_edge_id : act_node->output_edges())
+                ARM_COMPUTE_ERROR_ON(act_node->output(0) == nullptr || bn_node->output(0) == nullptr);
+
+                // Prevent fusion if batch normalization node has an output accessor
+                if(bn_node->output(0)->accessor() == nullptr)
                 {
-                    auto act_output_edge = g.edge(act_output_edge_id);
-                    if(act_output_edge != nullptr)
+                    // Get driving nodes of activation node
+                    std::vector<NodeIdxPair> act_driving_nodes;
+                    for(auto &act_output_edge_id : act_node->output_edges())
                     {
-                        ARM_COMPUTE_ERROR_ON(act_output_edge->consumer() == nullptr);
-                        act_driving_nodes.push_back({ act_output_edge->consumer_id(), act_output_edge->consumer_idx() });
+                        auto act_output_edge = g.edge(act_output_edge_id);
+                        if(act_output_edge != nullptr)
+                        {
+                            ARM_COMPUTE_ERROR_ON(act_output_edge->consumer() == nullptr);
+                            act_driving_nodes.push_back(
+                            { act_output_edge->consumer_id(), act_output_edge->consumer_idx() });
+                        }
                     }
-                }
 
-                // Set activation info to batch normalization
-                bn_node->set_fused_activation(act_node->activation_info());
+                    // Set activation info to batch normalization
+                    bn_node->set_fused_activation(act_node->activation_info());
 
-                // Remove activation node
-                g.remove_node(act_node->id());
+                    // Extract activation node accessor if any
+                    auto act_node_accessor = act_node->output(0)->extract_accessor();
 
-                // Update batch normalization node outputs
-                for(auto &driving_node : act_driving_nodes)
+                    // Remove activation node
+                    g.remove_node(act_node->id());
+
+                    // Update batch normalization node outputs
+                    for(auto &driving_node : act_driving_nodes)
+                    {
+                        g.add_connection(bn_node->id(), 0, driving_node.node_id, driving_node.index);
+                    }
+
+                    // Update accessor to batch normalization node
+                    bn_node->output(0)->set_accessor(std::move(act_node_accessor));
+                }
+                else
                 {
-                    g.add_connection(bn_node->id(), 0, driving_node.node_id, driving_node.index);
+                    ARM_COMPUTE_LOG_GRAPH_VERBOSE("Prevented fusion as batch normalization node has an output accessor\n");
                 }
             }
         }