Custom TF installation on top of NGC container

Hi, I’m having a problem running this notebook with TF NGC container version 20.01-tf2-py3:

I’m getting a segfault when doing inference:

output_dict = model(input_tensor)

But seems to be a TF problem, already adjusted at 2.0 nightly (

Point is, we do not have a 2.0 + version ready. What would be the most adequate way to them update the framework within this container but also keeping the GPU support?

I’ve tried a naive approach to update

pip install -U --pre tensorflow=="2.*")

But seems that inference is not running on GPU at all, only CPU.


The 20.02 NGC container should be released in the next few days and will include TF 2.1 with the needed fix.