From 6732ec3dfffebd8e17250803e27f1a4c5dedefae Mon Sep 17 00:00:00 2001 From: Akshay Modi Date: Mon, 9 Jul 2018 10:23:16 -0700 Subject: Skip calling back into python if only 1 gradient to aggregate PiperOrigin-RevId: 203786896 --- tensorflow/c/eager/tape.h | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) (limited to 'tensorflow/c') diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h index 734e712daa..1adb0458c3 100644 --- a/tensorflow/c/eager/tape.h +++ b/tensorflow/c/eager/tape.h @@ -520,7 +520,12 @@ Status GradientTape::ComputeGradient( } } else { any_gradient_nonzero = true; - auto new_gradients = vspace.AggregateGradients(grad_it->second); + Gradient* new_gradients = nullptr; + if (grad_it->second.size() == 1) { + new_gradients = grad_it->second.at(0); + } else { + new_gradients = vspace.AggregateGradients(grad_it->second); + } if (sources_set.find(grad_it->first) == sources_set.end()) { gradients.erase(grad_it); } else { -- cgit v1.2.3