PreTraining the model on image-net to reproduce results. causing errors after 1 epoc.

[muhammadumair894]

python train.py
--config /content/FasterViT/fastervit/configs/faster_vit_1_224_1k.yaml
--model faster_vit_1_224
--tag faster_vit_1_224_exp_1
--batch-size 64
--lr 0.005
--mesa 0.2
--model-ema
--opt adamw
--weight-decay 0.005
--amp
--input-size 3 224 224
--data_dir "/imagenet/"