Repository: open-mmlab/mmocr

Issue:
KIE type error exception
hi,
I am trying to run the KIE from the official colab notebook without success
the following code:

from mmocr.utils.ocr import MMOCR
mmocr = MMOCR(det='TextSnake', recog='SAR', kie='SDMGR')
mmocr.readtext('demo/demo_kie.jpeg', print_result=False, output='outputs/demo_kie_pred.jpg')

but im getting 
TypeError: 'int' object is not callable

code runs properly if i remove   kie='SDMGR'


see below trace

```
load checkpoint from http path: https://download.openmmlab.com/mmocr/textdet/textsnake/textsnake_r50_fpn_unet_1200e_ctw1500-27f65b64.pth
load checkpoint from http path: https://download.openmmlab.com/mmocr/textrecog/sar/sar_r31_parallel_decoder_academic-dba3a4a3.pth
/mmocr/mmocr/apis/inference.py:50: UserWarning: Class names are not saved in the checkpoint's meta data, use COCO classes by default.
  warnings.warn('Class names are not saved in the checkpoint\'s '
load checkpoint from http path: https://download.openmmlab.com/mmocr/kie/sdmgr/sdmgr_unet16_60e_wildreceipt_20210520-7489e6de.pth
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-30-bde23fb2c911>](https://localhost:8080/#) in <module>()
      1 from mmocr.utils.ocr import MMOCR
----> 2 mmocr = MMOCR(det='TextSnake', recog='SAR', kie='SDMGR')
      3 mmocr.readtext('demo/demo_kie.jpeg', print_result=False, output='outputs/demo_kie_pred.jpg')

9 frames
[/mmocr/mmocr/utils/ocr.py](https://localhost:8080/#) in __init__(self, det, det_config, det_ckpt, recog, recog_config, recog_ckpt, kie, kie_config, kie_ckpt, config_dir, device, **kwargs)
    394             self.kie_model = revert_sync_batchnorm(self.kie_model)
    395             self.kie_model.cfg = kie_cfg
--> 396             load_checkpoint(self.kie_model, kie_ckpt, map_location=self.device)
    397 
    398         # Attribute check

[/usr/local/lib/python3.7/dist-packages/mmcv/runner/checkpoint.py](https://localhost:8080/#) in load_checkpoint(model, filename, map_location, strict, logger, revise_keys)
    540         dict or OrderedDict: The loaded checkpoint.
    541     """
--> 542     checkpoint = _load_checkpoint(filename, map_location, logger)
    543     # OrderedDict is a subclass of dict
    544     if not isinstance(checkpoint, dict):

[/usr/local/lib/python3.7/dist-packages/mmcv/runner/checkpoint.py](https://localhost:8080/#) in _load_checkpoint(filename, map_location, logger)
    479            information, which depends on the checkpoint.
    480     """
--> 481     return CheckpointLoader.load_checkpoint(filename, map_location, logger)
    482 
    483 

[/usr/local/lib/python3.7/dist-packages/mmcv/runner/checkpoint.py](https://localhost:8080/#) in load_checkpoint(cls, filename, map_location, logger)
    249         mmcv.print_log(
    250             f'load checkpoint from {class_name[10:]} path: {filename}', logger)
--> 251         return checkpoint_loader(filename, map_location)
    252 
    253 

[/usr/local/lib/python3.7/dist-packages/mmcv/runner/checkpoint.py](https://localhost:8080/#) in load_from_http(filename, map_location, model_dir)
    288     if rank == 0:
    289         checkpoint = load_url(
--> 290             filename, model_dir=model_dir, map_location=map_location)
    291     if world_size > 1:
    292         torch.distributed.barrier()

[/usr/local/lib/python3.7/dist-packages/torch/hub.py](https://localhost:8080/#) in load_state_dict_from_url(url, model_dir, map_location, progress, check_hash, file_name)
    557     if _is_legacy_zip_format(cached_file):
    558         return _legacy_zip_load(cached_file, model_dir, map_location)
--> 559     return torch.load(cached_file, map_location=map_location)

[/usr/local/lib/python3.7/dist-packages/torch/serialization.py](https://localhost:8080/#) in load(f, map_location, pickle_module, **pickle_load_args)
    593                     return torch.jit.load(opened_file)
    594                 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
--> 595         return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
    596 
    597 

[/usr/local/lib/python3.7/dist-packages/torch/serialization.py](https://localhost:8080/#) in _legacy_load(f, map_location, pickle_module, **pickle_load_args)
    772     unpickler = pickle_module.Unpickler(f, **pickle_load_args)
    773     unpickler.persistent_load = persistent_load
--> 774     result = unpickler.load()
    775 
    776     deserialized_storage_keys = pickle_module.load(f, **pickle_load_args)

[/usr/local/lib/python3.7/dist-packages/torch/serialization.py](https://localhost:8080/#) in persistent_load(saved_id)
    728                 obj = data_type(size)
    729                 obj._torch_load_uninitialized = True
--> 730                 deserialized_objects[root_key] = restore_location(obj, location)
    731             storage = deserialized_objects[root_key]
    732             if view_metadata is not None:

[/usr/local/lib/python3.7/dist-packages/torch/serialization.py](https://localhost:8080/#) in restore_location(storage, location)
    815     else:
    816         def restore_location(storage, location):
--> 817             result = map_location(storage, location)
    818             if result is None:
    819                 result = default_restore_location(storage, location)

TypeError: 'int' object is not callable
```


Discussion:

Resolving patch:
diff --git a/mmocr/utils/ocr.py b/mmocr/utils/ocr.py
--- a/mmocr/utils/ocr.py
+++ b/mmocr/utils/ocr.py
@@ -329,8 +329,8 @@ def __init__(self,
         self.kie = kie
         self.device = device
         if self.device is None:
-            self.device = torch.cuda.current_device() \
-                if torch.cuda.is_available() else 'cpu'
+            self.device = torch.device(
+                'cuda' if torch.cuda.is_available() else 'cpu')
 
         # Check if the det/recog model choice is valid
         if self.td and self.td not in textdet_models:


Test patch:
diff --git a/tests/test_utils/test_ocr.py b/tests/test_utils/test_ocr.py
--- a/tests/test_utils/test_ocr.py
+++ b/tests/test_utils/test_ocr.py
@@ -81,6 +81,8 @@ def test_ocr_init(mock_loading, mock_config, mock_build_detector,
 
     def loadcheckpoint_assert(*args, **kwargs):
         assert args[1] == gt_ckpt[-1]
+        assert kwargs['map_location'] == torch.device(
+            'cuda' if torch.cuda.is_available() else 'cpu')
 
     mock_loading.side_effect = loadcheckpoint_assert
     with mock.patch('mmocr.utils.ocr.revert_sync_batchnorm'):
@@ -105,8 +107,7 @@ def loadcheckpoint_assert(*args, **kwargs):
             mock_config.assert_called_with(gt_cfg[-1])
             mock_build_detector.assert_called_once()
             mock_loading.assert_called_once()
-        device = torch.cuda.current_device() if \
-            torch.cuda.is_available() else 'cpu'
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         calls = [
             mock.call(gt_cfg[i], gt_ckpt[i], device=device) for i in i_range
         ]
@@ -167,8 +168,7 @@ def loadcheckpoint_assert(*args, **kwargs):
             mock_config.assert_called_with(gt_cfg[-1])
             mock_build_detector.assert_called_once()
             mock_loading.assert_called_once()
-        device = torch.cuda.current_device() if \
-            torch.cuda.is_available() else 'cpu'
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         calls = [
             mock.call(gt_cfg[i], gt_ckpt[i], device=device) for i in i_range
         ]