Cannot run the sample code ,"Check failed: oshape_sum == data_shape_sum "

https://colab.research.google.com/drive/13c3yapuhOjLiyCb_MsFSSrm2eSg6UWzP?usp=sharing