Trained neural network using DIGITS is not working in DRIVE PX2


I am using NVIDIA DIGITS (v6) for training a custom multi class object detector. I have used the prototxt file from NVIDIA github ( and changed the parameters for rmulti class object detection.

The training was successful and I can see two objects being recognized in DIGITS. I got the tensorRT bin file and I deployed the network in DRIVE PX2. I have modified the sample code in PX2 for multiple class object detection.

Apparently, while doing inference on the video stream, I notices only one class is being detected.

On debugging, I notices the the bbox/regressor layer in caffe model needs to have 4 times the number of classes as an output(Eg, if I have 2 classes, I should have 8 outputs). But changing the parameter in prototxt file before training gives an shape error in the next following layers.

Anyone knows the solution or how to implement the trained neural network in DRIVE PX2?

Thanks in advance