Solved by albanD in post #4 
              
              
                “b = (x==x.max(dim=1,keepdim=True)[0]).type(torch.FloatTensor)” is actually not differentiable and so that x.grad is none? 

If returns a Tensor that does not require gradients even though the input does. This means it is not differentiable yes. 

And is there a way to keep the maximum value to be …