Skip to content
Snippets Groups Projects
Unverified Commit a6b5e0e4 authored by agim-a's avatar agim-a Committed by GitHub
Browse files

[Fix] check for CLASSES in checkpoint meta (#4936) 

- check for CLASSES in checkpoint meta when key meta does not exists
parent 5ebba9a9
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
.dev_scripts/batch_test.py
+ 1
− 1
View file @ a6b5e0e4
...@@ -166,7 +166,7 @@ def main(): ...@@ -166,7 +166,7 @@ def main():
checkpoint = load_checkpoint(model, cpt, map_location='cpu') checkpoint = load_checkpoint(model, cpt, map_location='cpu')
# old versions did not save class info in checkpoints, # old versions did not save class info in checkpoints,
# this walkaround is for backward compatibility # this walkaround is for backward compatibility
if 'CLASSES' in checkpoint['meta']: if 'CLASSES' in checkpoint.get('meta', {}):
model.CLASSES = checkpoint['meta']['CLASSES'] model.CLASSES = checkpoint['meta']['CLASSES']
else: else:
model.CLASSES = dataset.CLASSES model.CLASSES = dataset.CLASSES
......
tools/analysis_tools/test_robustness.py
+ 1
− 1
View file @ a6b5e0e4
...@@ -294,7 +294,7 @@ def main(): ...@@ -294,7 +294,7 @@ def main():
model, args.checkpoint, map_location='cpu') model, args.checkpoint, map_location='cpu')
# old versions did not save class info in checkpoints, # old versions did not save class info in checkpoints,
# this walkaround is for backward compatibility # this walkaround is for backward compatibility
if 'CLASSES' in checkpoint['meta']: if 'CLASSES' in checkpoint.get('meta', {}):
model.CLASSES = checkpoint['meta']['CLASSES'] model.CLASSES = checkpoint['meta']['CLASSES']
else: else:
model.CLASSES = dataset.CLASSES model.CLASSES = dataset.CLASSES
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment