fixed `input_shape` error

minor updates for mlp_hook where error occurs when 'Y_hat' is not a tuple.