I don’t think that replacing the NaNs with other values would solve the issue, as I would assume that the backward pass would still create invalid gradients (you should verify it). The proper fix would be to avoid dividing by zero (or a very small number) in:
output_tensor = squared_norm * input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))

E.g. you could add a small eps value to the division.