Unable to reproduce results with CIFAR-10 dataset on efficientnet-l2-475

Hi,

I am running the following command to reproduce your results achieved with CIFAR-10 dataset on efficientnet-l2-475. I reduced the number of epochs just to see the script running from the very beginning till the end faster.
`python3 -m sam.sam_jax.train --dataset cifar10 --model_name efficientnet-l2-475 --output_dir /tmp/my_experiment --image_level_augmentations autoaugment --num_epochs 2 --sam_rho 0.05`

Nevertheless, I am getting the following error. How can I fix this?
`python3 -m sam.sam_jax.train --dataset cifar10 --model_name efficientnet-l2-475 --output_dir /tmp/my_experiment --image_level_augmentations autoaugment --num_epochs 2 --sam_rho 0.05`

Best regards,
Miguel