for visualization problems

[wefwefWEF2]

Thanks a lot for the great job!

I have some problems about visualization.
When I use
python visualize.py --config-file configs/cub.yaml OUTPUT_DIR /home/user/Downloads/TOAST-main/output MODEL.MODEL_ROOT /home/user/Downloads/TOAST-main/model/vit_b_bottom_up.pth MODEL.TEST_WEIGHT_PATH /home/user/Downloads/TOAST-main/output/CUB/vit_fb_ppt_small_patch16_224/lr0.01_wd0.0005/run4/test_CUB_logits_1.pth DATA.DATAPATH //home/user/Downloads/TOAST-main/visual_classification/src/data/datasets/CUB_200_2011/CUB_200_2011/

it shows AssertionError: Model type 'ssl-vit' not supported

But when I try to change type 'ssl-vit' to "vit_bottom_up" in cub.yaml, it works, but still can not load the model here, and shows key error

try:
model.load_state_dict(torch.load(cfg.MODEL.TEST_WEIGHT_PATH)["model"], strict=True)
except:
model.load_state_dict(torch.load(cfg.MODEL.TEST_WEIGHT_PATH)["model_state_dict"], strict=True)

and i check the keys in model, there are only dict_keys(['targets', 'joint_logits'])

How can I visualize it and load the model correctly?