MOCO V3 vit_small error: object has no attribute “num_tokens“

When I attempt to pre-train moco v3's vit_small model, I run into the following bug:

raise AttributeError("'{}' object has no attribute '{}'".format( AttributeError: 'VisionTransformerMoCo' object has no attribute 'num_tokens'

After changing the line
vits.py-line-66-LINK to
assert self.num_prefix_tokens == 1, 'Assuming one and only one token, [cls]' I don't see the bug anymore. It seems like the base class timm.models.vision_transformer has an argument named num_prefix_tokens but not num_tokens and hence vit_small is erroring out at the above mentioned line.

The command I used to run the code is:
python main_moco.py \ -a vit_small -b 1024 \ --optimizer=adamw --lr=1.5e-4 --weight-decay=.1 \ --epochs=400 --warmup-epochs=40 \ --stop-grad-conv1 --moco-m-cos --moco-t=.2 \ --dist-url 'tcp://localhost:8080' \ --multiprocessing-distributed --world-size 1 --rank 0 \ /data/

Please let me know if this is an accurate fix, or if I missed something. Thanks in advance!

基类vision_transformer中没有num_tokens，更改num_tokens为num_prefix_tokens，