Add test for BatchNorm running variables synchronization

tests/python/unittest/test_gluon.py

@@ -665,6 +665,34 @@ def transpose(shape):
         assert (layer(x).shape==ceil_out_shape)
 
 
+@with_seed()
+@pytest.mark.parametrize('cudnn_off', [True, False])
+@pytest.mark.parametrize('variable', ['running_var', 'running_mean'])
+def test_batchnorm_backward_synchronization(cudnn_off, variable):
+    """
+    Tests if synchronization of BatchNorm running variables is done correctly.
+    If not, the test sometimes fails - depending on the timing.
+    """
+    ctx = mx.cpu() if cudnn_off else mx.gpu()
+    read_op = 'layer.' + variable + '.data().asnumpy()'
+
+    for _ in range(20):
+        layer = nn.BatchNorm()
+        layer.initialize(ctx=ctx)
+        for _ in range(3):
+            data = mx.nd.random.normal(loc=10, scale=2, shape=(1, 3, 10, 10), ctx=ctx)
+            with mx.autograd.record():
+                out = layer(data)
+            out.backward()
+
+        # check if each read give the same value
+        var1 = eval(read_op)
+        for _ in range(10):
+            var2 = eval(read_op)
+            if (var1 != var2).any():
+                raise AssertionError("Two consecutive reads of " + variable + " give different results")
+
+
 @with_seed()
 def test_batchnorm():
     layer = nn.BatchNorm(in_channels=10)