def METHOD_NAME(self):
    """ Testing aggregation for torch model """
    aggregated_params = self.aggregator.aggregate(self.models, self.weights)
    # ===============================================================
    # Assert Federated Average
    # ===============================================================
    for key, val in aggregated_params.items():
        self.assertTrue(torch.isclose(val, self.model.state_dict()[key]).all())