FedProx regularization term
After a modification on self._model.init_params
the function __norm_l2
should be modified accordingly, otherwise in the for
loop, init_model
corresponds to the parameter's name and not to its value. This can be solved for instance as:
for current_model, init_model in zip(self.model().parameters(), self._model.init_params.values()):