Update train.py

vit_jax/train.py

@@ -72,7 +72,7 @@ def make_update_fn(*, apply_fn, accum_steps, tx):
 
 def train_and_evaluate(config: ml_collections.ConfigDict, workdir: str):
   """Runs training interleaved with evaluation."""
-
+  input_pipeline.get_datasets(config) # TODO：解决直接get_dataset_info加载信息失败的问题
   # Setup input pipeline
   dataset_info = input_pipeline.get_dataset_info(config.dataset, 'train')