Getting "OutOfMemory Error in GpuMemory: 0" from small CNN and small data-set

Hello, my objective is to train a very simple CNN on MNIST using Tensorflow, convert it to TensorRT, and use it to perform inference on the MNIST test set using TensorRT, all on a Jetson Nano, but I am getting several errors and warnings, including “OutOfMemory Error in GpuMemory: 0”. To try and reduce memory footprint, I tried also creating a script where I simply load the TensorRT model (that had already been converted and saved in the previous script) and use it to perform inference on a small subset of the MNIST test set (100 floating point values), but I am still getting the same out of memory error. The entire directory containing the TensorRT model is only 488 KB, and the 100 test points can’t be taking up very much memory, so I am confused about why GPU memory is running out. What could be the reason for this, and how can I solve it?

I am attaching the console output from the 2 Python scripts below as text files. These text files and the Python scripts which generated them can be found on this Gist:
z_console_output.txt (39.9 KB) zz2_trt_infer_console_output.txt (21.8 KB)

Another thing which seems suspicious is that some of the Tensorflow logging info messages are being printed multiple times, EG “Successfully opened dynamic library libcudart”, “Successfully opened dynamic library libcublas”, “ARM64 does not support NUMA - returning NUMA node zero”. What could be the reason for this (EG dynamic libraries being opened over and over again), and could this have something to do with why the GPU memory keeps running out?


Just check your log and it looks like everything works fine. (no error message).
Please let me know if anything missing.

Here is our end-to-end sample for MNIST for your reference: