Yes it seems to work quite nicely 
values = torch.tensor([0.20, 0.40, 0.60, 0.80], requires_grad=True)

def gradual (values, p):
    return torch.where (values == values.max(), 1.0 - (1.0 - values)**p, values**p)


softmax_values = torch.stack([gradual(values, p) for p in np.arange(1, 20)])

for input, outputs in zip(values, softmax_values.T):
    c = plt.gca()._get_lines.get_next_color()
    plt.plot([0], [input], 'o', c=c)
    plt.plot(outputs.detach(), c=c)
    
plt.xlabel("p")
plt.show()

loss = softmax_values.mean()
loss.backward()
values.grad


tensor([0.0206, 0.0365, 0.0822, 0.0206])
Do you know if that is the case generally, that using torch.where(x == max(x), t1, t2) is differentiable while torch.argmax(x) is not? I guess they have different use-cases but this feels like a very nice “trick” to maintain gradient.