Download Imagenet for pytorch fails

[Tobias-Buerger]

Downloading Imagenet for pytorch framework fails.

Description

When running the dataset_setup.py for pytorch we get a NameError caused by a copy&paste bug from the jax download section.

Steps to Reproduce

Running

python datasets/dataset_setup.py --data_dir $DATA_DIR --temp_dir $DATA_DIR/tmp --imagenet --imagenet_train_url <url> --imagenet_val_url <url> --framework pytorch

Source or Possible Fix

Change variable imagenet_jax_data_dir to imagenet_pytorch_data_dir in dataset_setup.py:522 and dataset_setup.py:524

[priyakasimbeg]

Thanks for filing this issue. I happened to be working on a fix, just merged it.

[priyakasimbeg]

Reopened this because I'm getting an "input format not supported by decoder error" at the extract step.

[priyakasimbeg]

Submitted fix for reading the tar file.