Hi, Thanks! but what you proposed, won’t work. By default, if the grad is None, that’s ok for backprop on a single machine. But then, I am doing distributed training, and I need to average out the gradients over multiple workers. When the workers encounter None, they give error.