Investigate Error "f capturable=True, params and state_steps must be CUDA or XLA tensors" when nn module inlining enabled.

[laithsakka]

REPO:
enable nn module inlinining then run
TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1 PYTORCH_TEST_WITH_DYNAMO=1 python test/test_quantization.py -k test_qat_embeddingbag_linear

PYTORCH_TEST_WITH_DYNAMO=1 TORCHDYNAMO_INLINE_INBUILT_NN_MODULES=1 python test/test_quantization.py -k test_resnet_base

nine total failures with this message, check: https://hud.pytorch.org/pr/125354
cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @vincentqb @janeyx99 @crcrpar @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

[laithsakka]

note: those only fails when dynamo is enabled on the test

[janeyx99]

What is module inlining/how does this affect the optimizer?

This looks like an optimizer error that happens when capturable is flipped on but the state is not ready to be captured.

[laithsakka]

I had offline chat with @mlazos he is already working on this or something very close to this, he said he will take over

[laithsakka]

@mlazos @anijain2305 this is fixed with recent Lazos changes

[laithsakka]

fixed by #125826