def test_checkpoint_cb_for_save_op(): """Test checkpoint cb for save op.""" parameter_list = [] one_param = {} one_param['name'] = "conv1.weight" one_param['data'] = Tensor(np.random.randint(0, 255, [1, 3, 224, 224]), dtype=mstype.float32) parameter_list.append(one_param) _checkpoint_cb_for_save_op(parameter_list)
def test_checkpoint_cb_for_save_op_update_net(): """Test checkpoint cb for save op.""" parameter_list = [] one_param = {} one_param['name'] = "conv.weight" one_param['data'] = Tensor(np.ones(shape=(64, 3, 3, 3)), dtype=mstype.float32) parameter_list.append(one_param) net = Net() _set_cur_net(net) _checkpoint_cb_for_save_op(parameter_list) assert net.conv.weight.data.asnumpy()[0][0][0][0] == 1