As @danielgrittner pointed out, some models can have different naming conventions but follow the same pattern, this pr fixes the issue by using the base_model_prefix and finding the transformer block name, instead of always using transformer_h_X.
Tests work the same as they did, as this doesn't affect the workings of the partitioning, we just are generalizing the way to identify transformer blocks.
As @danielgrittner pointed out, some models can have different naming conventions but follow the same pattern, this pr fixes the issue by using the base_model_prefix and finding the transformer block name, instead of always using transformer_h_X.
Tests work the same as they did, as this doesn't affect the workings of the partitioning, we just are generalizing the way to identify transformer blocks.