change model output to tensor

Currently output is a list of Tensors. It would be nicer/more convenient to have a tensor with one extra dimension.