It takes ages to load on every step even though my dataset is quite small.