So FYI if someone runs into the same issue.
Increasing the training epochs of the DIETClassifier fixed the issue. For dev purposes, I decreased it to 20 epochs to speed up training times. That was not a good idea 
.
From my experience, the threshold is around a min of 60 epochs. A lower value decreases the entity recognition significantly.