[Lazy init] Fix edge cases
This PR fixes the flaky circle ci regarding tests such as `tests/test_modeling_xlnet.py::XLNetModelTest::test_save_load_fast_init_to_base`. 

Luckily, the test caught those edge cases ;-) 