MKLDNN fallback when not recording gradients and calling backwards

src/operator/nn/mkldnn/mkldnn_convolution.cc

@@ -597,7 +597,7 @@ void MKLDNNConvolutionBackward(const nnvm::NodeAttrs& attrs, const OpContext &ct
     out_grad = out_grad.Reorder2Default();
 
   mkldnn::convolution_forward::primitive_desc fwd_pd = GetConvFwdImpl(
-      full_param, ctx.is_train, data, weight, bias, out_grad);
+      full_param, ctx.need_grad, data, weight, bias, out_grad);
   const ConvolutionParam &param = full_param.conv_param;
 
   CHECK_NE(req[conv::kWeight], kWriteInplace) << "cannot write weight inplace";

tests/python/unittest/test_gluon.py

@@ -1209,7 +1209,7 @@ def test_zero_grad():
     grad = net.collect_params()['test_zero_grad_weight'].grad()
     assert_almost_equal(grad.asnumpy(), grad.asnumpy() * 0)
 
-def check_hybrid_static_memory(**kwargs):
+def check_hybrid_static_memory(train_modes, **kwargs):
     x = mx.nd.random.uniform(shape=(2, 3, 32, 32))
     x.attach_grad()
 
@@ -1221,27 +1221,29 @@ def check_hybrid_static_memory(**kwargs):
     net1(x)
     net2(x)
 
-    def test(net, x):
-        with mx.autograd.record():
+    def test(net, x, record=True):
+        with mx.autograd.record(record):
             y = net(x) + net(x)
             y.backward()
 
         grads = {k: v.grad() for k, v in net.collect_params().items() if v.grad_req != 'null'}
 
         return y, grads
 
-    y1, grads1 = test(net1, x)
-    y2, grads2 = test(net2, x)
+    for record in train_modes:
+        y1, grads1 = test(net1, x, record)
+        y2, grads2 = test(net2, x, record)
 
-    assert_almost_equal(y1.asnumpy(), y2.asnumpy(), rtol=1e-3, atol=1e-5)
-    for key in grads1:
-        assert_almost_equal(grads1[key].asnumpy(), grads2[key].asnumpy(), rtol=1e-3, atol=1e-5)
+        assert_almost_equal(y1.asnumpy(), y2.asnumpy(), rtol=1e-3, atol=1e-5)
+        for key in grads1:
+            assert_almost_equal(grads1[key].asnumpy(), grads2[key].asnumpy(), rtol=1e-3, atol=1e-5)
 
 @with_seed()
 def test_hybrid_static_memory():
-    check_hybrid_static_memory()
-    check_hybrid_static_memory(static_alloc=True)
-    check_hybrid_static_memory(static_alloc=True, static_shape=True)
+    check_hybrid_static_memory(train_mode=[True, False])
+    check_hybrid_static_memory(train_mode=[True, False], static_alloc=True)
+    # TODO: MKLDNN (issue #13445) does not work with static_shape backwards
+    check_hybrid_static_memory(train_mode=[True], static_alloc=True, static_shape=True)
 
 def check_hybrid_static_memory_switching(**kwargs):
     net = gluon.model_zoo.vision.get_resnet(