Networks with BatchNorm and gradients

[hskramer]

I have set the i.grad != "null" which I found in here but now I receive error conv0_weight' was not initialized on context cpu(0). It was only initialized on [gpu(0), gpu(1)], but if I run it on just one gpu I receive the assertion error len(grads) == len(param_names) this occurs on all nets with a BatchNorm ones I have written and one in gluon's model zoo tried Inception, Densenet121 error:

grads = [i.grad() for i in net.collect_params().values() if i.grad_req != 'null']
50 assert len(grads) == len(param_names)
51

in (.0)
47 sw.add_graph(net)
48
---> 49 grads = [i.grad() for i in net.collect_params().values() if i.grad_req != 'null']
50 assert len(grads) == len(param_names)
51

/usr/local/lib/python3.6/dist-packages/mxnet/gluon/parameter.py in grad(self, ctx)
570 "Cannot get gradient array for Parameter '%s' "
571 "because grad_req='null'"%(self.name))
--> 572 return self._check_and_get(self._grad, ctx)
573
574 def list_grad(self):

/usr/local/lib/python3.6/dist-packages/mxnet/gluon/parameter.py in _check_and_get(self, arr_list, ctx)
225 "Parameter '%s' was not initialized on context %s. "
226 "It was only initialized on %s."%(
--> 227 self.name, str(ctx), str(self._ctx_list)))
228 if self._deferred_init:
229 raise DeferredInitializationError(

RuntimeError: Parameter 'densenet0_conv0_weight' was not initialized on context cpu(0). It was only initialized on [gpu(0), gpu(1)].