Move JK check to on-demand (#121182)
Summary: Some tests are failing due to checking JK during forking. Lets move the JK check to on-demand.
Differential Revision: D54518293
Pull Request resolved: https://github.com/pytorch/pytorch/pull/121182
Approved by: https://github.com/aakhundov
