Add correctness check for sharded checkpoint test

[mrwyattii]

Discussion on #2379 has indicated that there are correctness issues when loading certain models from sharded checkpoints.

Should be merged after #2662

@RezaYazdaniAminabadi