Convert the max_mem_mb to bytes before comparing with size of dataset
Get the size of the dataset in memory rather than the number of datapoints