model_cls.emb_dim_seg not taken into account #15

Hi,

I noticed that at inference time for the segmentation task, the script with my custom model (specified with model_cls) was failing because the shape of the best checkpoint and the dimension of the embeddings were not matching. Indeed, my custom model has feature dimension 1536 for linear probing but 768 for segmentation. Checkpoint was expecting 768 but found 1536. I've added the following in train_eval_probe.py:359

    if model_cls is not None:
        OmegaConf.set_struct(cfg, False)
        cfg.pretrained_model = {"emb_dim": model_cls.emb_dim}
        if hasattr(model_cls, "emb_dim_seg"):
            cfg.pretrained_model.update({"emb_dim_seg": model_cls.emb_dim_seg})