ValueError: expected 4D input (got 3D input)