Problem with create model from pretrained checkpoint

[badbul]

Hi! Thanks for your repo, it's very useful.
But i met with some problem when tried to create model from pretrained checkpoint.
I'm using effdet==0.3.0 package. When i run the folowing code:
model = create_model("tf_efficientdet_d5", bench_task="train", num_classes=3, image_size=(512, 512), bench_labeler=True, pretrained=False, checkpoint_path="tf_efficientdet_d5_51-c79f9be6.pth" )
I get the following error:
'size mismatch for class_net.predict.conv_pw.weight' on line 'load_checkpoint(model, checkpoint_path, use_ema=checkpoint_ema)' of the file 'effdet/factory.py"'

After that in the file 'effdet/factory.py' i placed code :
' if checkpoint_path:
load_checkpoint(model, checkpoint_path, use_ema=checkpoint_ema)'
before code:
if num_classes is not None and num_classes != config.num_classes: model.reset_head(num_classes=num_classes)
and error is gone.
Did you met with such problem?