Performance issue in baselines/cifar/ (by P3)

Hello! I've found a performance issue in utils.py /: `dataset.batch(batch_size, drop_remainder=drop_remainder)`[(here)](https://github.com/google/uncertainty-baselines/blob/56434c78f96c488d7839a3bde61db1fd2ce2901a/baselines/cifar/utils.py#L117) should be called before `dataset.map(preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)`[(here)](https://github.com/google/uncertainty-baselines/blob/56434c78f96c488d7839a3bde61db1fd2ce2901a/baselines/cifar/utils.py#L115), which could make your program more efficient.

Here is [the tensorflow document](https://tensorflow.google.cn/guide/data_performance?hl=zh_cn#vectorized_mapping) to support it.

Besides, you need to check the function `preprocess` called in `dataset.map(preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)` whether to be affected or not to make the changed code work properly. For example, if `preprocess` needs data with shape (x, y, z) as its input before fix, it would require data with shape (batch_size, x, y, z) after fix.

Looking forward to your reply. Btw, I am very glad to create a PR to fix it if you are too busy.