From 7b4b94fd93afe389bcaa3b38e09abe0dc2a24be8 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 13 Jun 2022 16:25:17 +0800
Subject: [PATCH] fix(imperative): fix the segmentfault when reduce backward

GitOrigin-RevId: 8a3e63d4f538ace0e06b3df6aaaa080633a3e525
---
 imperative/python/src/grad_override.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/imperative/python/src/grad_override.cpp b/imperative/python/src/grad_override.cpp
index be6c95f9..c576aab2 100644
--- a/imperative/python/src/grad_override.cpp
+++ b/imperative/python/src/grad_override.cpp
@@ -408,7 +408,7 @@ std::optional<ValueRefList> reduce_grad_rule(
             [shapes = std::move(input_shapes), axis, keepdim](Span<ValueRef> grads) {
                 mgb_assert(grads.size() == 1);
                 ValueRef grad = grads[0];
-                if (!keepdim) {
+                if (!keepdim && grad) {
                     auto&& grad_op = AddAxis::make(std::vector<int32_t>({axis}));
                     grad = imperative::apply(*grad_op, grad)[0];
                 }