Repository: pytorch
Issue: `torch.nn.KLDivLoss` fails `test_grad` and `test_gradgrad` 
Labels: module: nn,triaged,needs research

### 🐛 Describe the bug

As title, discovered in https://github.com/pytorch/pytorch/pull/115584 and seems to occur for all ModuleInputs

e.g. 

```
======================================================================
ERROR: test_grad_nn_KLDivLoss_cuda_float64 (__main__.TestModuleCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_utils.py", line 2652, in wrapper
    method(*args, **kwargs)
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_utils.py", line 2652, in wrapper
    method(*args, **kwargs)
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_device_type.py", line 429, in instantiated_test
    raise rte
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_device_type.py", line 416, in instantiated_test
    result = test(self, **param_kwargs)
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_utils.py", line 1540, in wrapper
    fn(*args, **kwargs)
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_modules.py", line 122, in test_wrapper
    return test(*args, **kwargs)
  File "/data/users/mg1998/pytorch/test/test_modules.py", line 506, in test_grad
    self._test_gradients_helper(device, dtype, module_info, training, gradcheck)
  File "/data/users/mg1998/pytorch/test/test_modules.py", line 478, in _test_gradients_helper
    self.assertTrue(check(fn_to_gradcheck, flat_input, nondet_tol=gradcheck_nondet_tol))
  File "/data/users/mg1998/pytorch/torch/testing/_internal/common_utils.py", line 4422, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/data/users/mg1998/pytorch/torch/autograd/gradcheck.py", line 2070, in gradcheck
    return _gradcheck_helper(**args)
  File "/data/users/mg1998/pytorch/torch/autograd/gradcheck.py", line 2099, in _gradcheck_helper
    _gradcheck_real_imag(
  File "/data/users/mg1998/pytorch/torch/autograd/gradcheck.py", line 1487, in _gradcheck_real_imag
    gradcheck_fn(
  File "/data/users/mg1998/pytorch/torch/autograd/gradcheck.py", line 1920, in _fast_gradcheck
    _check_analytical_numerical_equal(
  File "/data/users/mg1998/pytorch/torch/autograd/gradcheck.py", line 1849, in _check_analytical_numerical_equal
    raise GradcheckError(
torch.autograd.gradcheck.GradcheckError: Jacobian mismatch for output 0 with respect to input 0,
numerical:tensor(nan, device='cuda:0', dtype=torch.float64)
analytical:tensor(nan, device='cuda:0', dtype=torch.float64)

The above quantities relating the numerical and analytical jacobians are computed 
in fast mode. See: https://github.com/pytorch/pytorch/issues/53876 for more background 
about fast mode. Below, we recompute numerical and analytical jacobians in slow mode:

Numerical:
 tensor([[nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan],
        [nan]], device='cuda:0', dtype=torch.float64)
Analytical:
tensor([[    nan],
        [    nan],
        [-2.1202],
        [    nan],
        [-2.0278],
        [    nan],
        [    nan],
        [-1.1458],
        [    nan],
        [    nan],
        [    nan],
        [-1.9755],
        [    nan],
        [    nan],
        [    nan],
        [-1.3503],
        [    nan],
        [    nan],
        [-2.1348],
        [    nan],
        [    nan],
        [-0.4169],
        [ 0.5499],
        [    nan],
        [ 0.3591],
        [-1.5428],
        [-1.9322],
        [    nan],
        [-1.0274],
        [    nan],
        [-2.1358],
        [    nan],
        [    nan],
        [-1.8342],
        [-1.9648],
        [    nan],
        [-1.7928],
        [    nan],
        [    nan],
        [    nan],
        [-1.7968],
        [-1.6731],
        [-0.0203],
        [-0.8002],
        [-1.9048],
        [    nan],
        [    nan],
        [-0.3072],
        [-1.5951],
        [-0.7291],
        [-1.4519],
        [    nan],
        [    nan],
        [    nan],
        [    nan],
        [    nan],
        [-1.9169],
        [    nan],
        [-1.4493],
        [    nan],
        [-0.6638],
        [    nan],
        [    nan],
        [-1.6942],
        [-2.1172],
        [    nan],
        [    nan],
        [    nan],
        [    nan],
        [    nan],
        [    nan],
        [    nan],
        [-0.6606],
        [-1.5360],
        [    nan],
        [-2.1227],
        [-1.7354],
        [    nan],
        [-1.0844],
        [    nan],
        [    nan],
        [    nan],
        [-1.6773],
        [-1.8423],
        [    nan],
        [-1.9218],
        [    nan],
        [ 0.4764],
        [    nan],
        [    nan],
        [-1.1524],
        [ 0.0306],
        [    nan],
        [    nan],
        [-1.8915],
        [-0.2729],
        [    nan],
        [    nan],
        [-1.3643],
        [    nan]], device='cuda:0', dtype=torch.float64)

The max per-element difference (slow mode) is: nan.


To execute this test, run the following from the base repo dir:
     python test/test_modules.py -k test_grad_nn_KLDivLoss_cuda_float64

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0
```

Needs to be root caused.

### Versions

main

cc @albanD @mruberry @jbschlosser @walterddr