Update optimizer research codes usage of `tff.learning` to reflect new return value structure.
@@ -55,6 +55,11 @@ class AdaptiveFedAvgTest(tf.test.TestCase):
state = iterative_process.initialize()
for round_num in range(num_rounds):
iteration_result = iterative_process.next(state, client_datasets)
+ if hasattr(iteration_result.metrics, 'train'):
+ # tff.learning returns a nested tuple of metrics, we only compare
+ # against `train`.
+ train_outputs.append(iteration_result.metrics.train)
+ else:
train_outputs.append(iteration_result.metrics)
logging.info('Round %d: %s', round_num, iteration_result.metrics)
logging.info('Model: %s', iteration_result.state.model)
