train_test_split: divide classes given class in other column

Question

I have this pandas dataframe:

ID label age   sex  type    
0   0    20.0   0   0.0 
1   1    20.0   0   1.0 
2   0    20.0   1   0.0 
3   1    20.0   1   1.0