I'm experimenting with automatic cleanup for eager tensors, like JavaCPP does for tensors. I have it working on CPU using JavaCPP's detection (it allocates over the memory limit during execution, then pulls back). However, on GPU I'm forced to catch OOM exceptions on execute, then run cleanup and re-try. The catching works fine, but when I execute the same op again I get "expected 2 inputs, got 0". This is despite TFE_OpGetInputLength(
opHandle, "value", TF_Status.newStatus())
returning 1 for both arguments.
To see this, run something that would force an OOM on a GPU.
Any thoughts on how to get around this? The easiest option I can see is to store set inputs/attributes in the builder and re-build the native op each call, which is less than ideal.