Hi @albanD,
a solution to the problem would simply be to make a copy of the whole model, not only of the state_dict and operate on the copied model:
model_copy = copy.deepcopy(model)

        logits = model_copy(x)
        return objective(logits, y)

To be honest, this was not my idea.
Anyway, thanks for your help!
Jürgen