TF keras.py _wrap_tape_gradient breaks for arrays

[arewellborn]

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.