“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 1 and the others to be 0 in a tensor but the tensor still could have gradient?

If you think about gradient as “how much the output change if I change this input a little”, then you can see that if the output of your function is “the maximum value to be 1 and the others to be 0”, then all the gradients will be 0. Because when you change any input a little (if the maximum one does not change), then all the outputs will remain the same.