class_data, counts = get_data_distribution([train_dataset,valid_dataset,test_dataset])
visualize_data([class_names[cd] for cd in class_data],counts.numpy())

        __
        __