[Compile] Reduce guards and recompiles for TensorDict under torch.compile

[Compile] Reduce guards and recompiles for TensorDict under torch.compile #2023