import numpy as np
import mindspore
from mindspore import Tensor, save_checkpoint, load_checkpoint

if __name__ == "__main__":
    data = Tensor(np.random.uniform(0, 255, (1, 2)), mindspore.uint8)
    save_obj = [{
        "name": "test_data",
        "data": data
    }]

    test_save_path = "test_save_load.ckpt"
    save_checkpoint(save_obj, test_save_path)

    param_dict = load_checkpoint(test_save_path)
    print(param_dict["test_data"].dtype)

[ERROR] ME(38583:139942186227520,MainProcess):2021-11-08-12:08:58.359.733 [mindspore/train/serialization.py:467] Failed to load the checkpoint file `XXX/test_save_load.ckpt`.
Traceback (most recent call last):
  File "XXX/mindspore/train/serialization.py", line 440, in load_checkpoint
    np_type = tensor_to_np_type[data_type]
KeyError: 'UInt8'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "XXX/test_save_load_ckpt.py", line 15, in <module>
    param_dict = load_checkpoint(test_save_path)
  File "XXX/mindspore/train/serialization.py", line 468, in load_checkpoint
    raise RuntimeError(e.__str__())
RuntimeError: 'UInt8'