can this be updated to tf2?

[constant5]

[jperl]

Hi @constant5 I do not have time to do that right now but I would be happy to review a PR.

[constant5]

OK, thanks. I got it working by revising the code using compat.v1 where needed, but if I get a chance I will try and implement in Keras so it is a little more portable.