Skip to content
Snippets Groups Projects
Commit a6857292 authored by Vladislav Sovrasov's avatar Vladislav Sovrasov Committed by Kai Chen
Browse files

Fix acces to a wrong key in checkpoint (#734)

parent 0d9dbcc4
No related branches found
No related tags found
No related merge requests found
......@@ -34,7 +34,7 @@ def init_detector(config, checkpoint=None, device='cuda:0'):
if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint)
if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['classes']
model.CLASSES = checkpoint['meta']['CLASSES']
else:
warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use COCO classes by default.')
......
......@@ -79,7 +79,7 @@ def main():
cfg.checkpoint_config.meta = dict(
mmdet_version=__version__,
config=cfg.text,
classes=train_dataset.CLASSES)
CLASSES=train_dataset.CLASSES)
# add an attribute for visualization convenience
model.CLASSES = train_dataset.CLASSES
train_detector(
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment