I think I found a bug for this docmenthttps://pytorch.org/tutorials/advanced/torch_script_custom_ops.html, that in the code
torch::Tensor output = torch::from_blob(output.ptr<float>(), /*sizes=*/{8, 8});

it should be
torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{8, 8});