AssertionError

[feufhd]

Thanks for your work!
I want to explore diffusion models based on U-DiT, and I met errors during training (The command is: accelerate launch --multi_gpu --main_process_port 12759 --num_processes 16 --mixed_precision fp16 train_accelerate.py --data-path /imagenet/train/ --feature-path=/U-DiT/in1k_features/imagenet_feature --image-size=256 --model=U-DiT-B --epochs=80).

[rank0]: Traceback (most recent call last):
[rank0]: File "/U-DiT/train_accelerate.py", line 303, in
[rank0]: main(args, unparsed)
[rank0]: File "/U-DiT/train_accelerate.py", line 189, in main
[rank0]: dataset = CustomDataset(args.feature_path)
[rank0]: File "/U-DiT/utils/fastdit_kit.py", line 38, in init
[rank0]: assert len(self.features_files) == len(self.features_files) == 1281167 # ImageNet
[rank0]: AssertionError

I'm pretty sure that I've downloaded your given imagenet_feature.tar and unzipped it successfully. I print len(self.features_files) and it is 40031, which is 1/32 of the length of IN1K (1281167). Is it OK to delete the assertion here? Or what else did I neglect?

[YuchuanTian]

Before deleting this assertion, just make sure that the total number of features is 1281167. Just concat all features on the batch dimension and check.