fix Unflatten out_returnn_axis_from_torch_axis

[vieting]

The elif branch that I edited is supposed to be the case where i refers to a dim which outside (larger than) the unflattened dims. Equality should also be contained in this case.