diff --git a/demo/image_demo.py b/demo/image_demo.py index 87d6d6c41..86396799f 100644 --- a/demo/image_demo.py +++ b/demo/image_demo.py @@ -1,6 +1,8 @@ # Copyright (c) OpenMMLab. All rights reserved. from argparse import ArgumentParser +from mmcv.cnn.utils.sync_bn import revert_sync_batchnorm + from mmseg.apis import inference_segmentor, init_segmentor, show_result_pyplot from mmseg.core.evaluation import get_palette @@ -26,6 +28,8 @@ def main(): # build the model from a config file and a checkpoint file model = init_segmentor(args.config, args.checkpoint, device=args.device) + if args.device == 'cpu': + model = revert_sync_batchnorm(model) # test a single image result = inference_segmentor(model, args.img) # show the results