Torchvision Faster R-CNN custom

I’m currently implementing the model introduced in here, where it just simply adds a branch that detects rotated bounding boxes to Faster R-CNN roi head. I’ve adapted torchvision Faster R-CNN model at GitHub - crazyboy9103/oriented_detection, but it does not seem to work well. It is able to detect bounding boxes, but their classifications are mostly wrong and I don’t know where to fix. The dataset I’m using is MVTec rotated screws dataset.

Could anyone help me with this matter?
In the repository, I’ve mainly adapted models/detection/,,