Drug data split

[HITzhongyu]

Hi，
I try to use the main_geom_drugs,py to run , but it seems to have some error,

And I also try to solve it, but it maybe is the build_geom_dataset.py and line 101 data_list = [data_list[i] for i in perm], this problem is because the data_list contains subarrays of varying shapes,and in line 107 np.spilt need same shape,
So how can I solve this problem?

best regard,
Zhongyu