xla

50141f2e - [FSDP] Fix pre-hook check when module is only used in forward (#4670)

[FSDP] Fix pre-hook check when module is only used in forward (#4670) In a corner case, a sub-module is only used in forward but not in backward. The assert_state will fail because the module is still in BACKWARD_IDLE state. Add condition to make sure the module has run in backward.