use smaller shape to prevent DebugMode to return error sometimes due to different rounding on CPU and GPU with the Sum op.