Solved by lhoestq in post #2 
              
              
                Hi ! 
There are two optimizations that you can use: 

Use the Array3D feature type. Otherwise it will consider your arrays to be lists of arbitrary sizes, and it takes some time to collate them again into one numpy array.

features = Features({
    **train_ds.features,
    "image": Array3D(dtype="in…