Error with unnormalized quantum state

[Nikait]

Hello! I've been learning your code (model.py) and noticed a bug in the function below:

def apply_unitaries(base_states, unitary_params, enc, qdev, n_qbs, coeffs):
   repeated_base = base_states.repeat(1, unitary_params.shape[1]).view(-1, 2 ** n_qbs)
   qdev.set_states(repeated_base)

   enc(qdev, unitary_params.view(-1, unitary_params.shape[-1]))
   states = qdev.get_states_1d().view(*unitary_params.shape[:2], 2 ** n_qbs)
   lcs = torch.einsum('bwi,bw->bi', states, coeffs)
   return lcs

After an iteration of QSVT the amplitude vectors of base_states aren't normalized:
For all i=0...repeated_base.shape[0]-1: torch.linalg.vector_norm(repeated_base[i]) != 1.

But the amplitude vector of any quantum state must be normalized...