Describe the bug:

keypoints_inputs() method does not properly update, and instead runs on the previous iteration of keypoints constructed during forward pass.

To Reproduce
Steps to reproduce the behavior:

1. Create a numerical calibrator with learned input keypoints
2. Train the module on a few examples
3. call calibrator.keypoints_inputs()
4. Return value will be from the values from before the final training step.

Expected behavior
The final training step updates should be properly referenced. When using learned input keypoints, the keypoints_inputs function should generate the input keypoints directly from the logits.

Additional context
N/A