Solved by albanD in post #6 
              
              
                The y should be a single Tensor right? 
So flat_y here should be a 1D Tensor. And flat_y[i] will be a single number. So it is a scalar output. 
Is the error coming from this function?