use_ml_decoder_head

pull/1012/head
talrid 3 years ago
parent 41559247e9
commit b6c180ef19

@ -381,7 +381,7 @@ def main():
bn_eps=args.bn_eps, bn_eps=args.bn_eps,
scriptable=args.torchscript, scriptable=args.torchscript,
checkpoint_path=args.initial_checkpoint, checkpoint_path=args.initial_checkpoint,
ml_decoder_head=args.use_ml_decoder_head) use_ml_decoder_head=args.use_ml_decoder_head)
if args.num_classes is None: if args.num_classes is None:
assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.' assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'
args.num_classes = model.num_classes # FIXME handle model default vs config num_classes more elegantly args.num_classes = model.num_classes # FIXME handle model default vs config num_classes more elegantly

Loading…
Cancel
Save