Hi!

I'm trying to use a Pytorch Tabular Model inside a Federated Learning environment, so I need to obtain the weights from different nodes, aggregate them and load the new weights inside each node.

The extracting part is OK, but i am having trouble in the loading part. After I have defined my model and train it in each model with:

node_x_model = TabularModel(data_config=data_config, model_config=model_config, optimizer_config=OptimizerConfig(), trainer_config=trainer_config, experiment_config=experiment_config)
node_x_model.fit(train=trainset, seed = 42)

I am able to extract to weights correctly. After operate with these weights, I obtain a new weights that we'll call new_model, that i load with:
node_x_model.model.load_state_dict(new_model)

Up to this point, there is no problem, but when the next round start and each node has to train again, using the same sentence:

node_x_model.fit(train=trainset, seed = 42)

Each node obtained results are the same that the corresponding node obtained the last round. It seems that the function fit is overwriting the loaded weights. I have revised the API docs, but I have not found the solution to this. Does anyone know how to do this?

Thanks in advance.