AttributeError: 'list' object has no attribute 'argmax' and numpy.AxisError: axis 2 is out of bounds for array of dimension 1 #335