[torch.compile] Bunch of small changes needed for enabling torch.compile

[torch.compile] Bunch of small changes needed for enabling torch.compile #10297