TF keras.py _wrap_tape_gradient breaks for arrays

I am passing in batches into tf.GradientTape() loop and calling my model for a prediction. When I call tape.gradient(prediction, batch), smdebug throws "ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()" for line 1209 at the file below.

https://github.com/awslabs/sagemaker-debugger/blob/master/smdebug/tensorflow/keras.py#L1209

Also, I am using tensorflow version 2.1. I am not sure if any updates can make it into the past versions (it looks like the AWS image I have started from used v0.9.3).

Currently I have uninstalled smdebug and it works as expected.

Start at smdebug/tensorflow/keras.py line 1209 and reproduce the failure with TensorFlow 2.1, tf.GradientTape(), and array batches. Confirm that calling tape.gradient(prediction, batch) no longer raises the ambiguous truth-value error while smdebug is installed.

Written by the indexing model from the issue text.