def aggregate_key(self, aggregate_key):
        """ Aggregate values from key and put them into the top-level dictionary """
        aggregation = self.data_dict[aggregate_key]  # List of dictionaries of numpy arrays/scalars

        # Aggregate sub batch data
        data_dict_keys = {y for x in aggregation for y in x.keys()}

        for key in data_dict_keys:
            # Just average all the statistics from the loss function
            stacked = np.stack([d[key] for d in aggregation], axis=0)
            self.data_dict[key] = np.mean(stacked, axis=0)
Aggregate values from key and put them into the top-level dictionary