NTS-Net Implementation in Fast.AI

Hi all, I want to try to implement NTS-Net (https://arxiv.org/abs/1809.00287) as it had good results on the UCSD Birds-200-2011 dataset and want to try to bring it into Fast.AI. This model consists of three parts, a Navigator, Teacher, and Scrutinizer (hence NTS) built on-top of a resnet50 model. However, when I try to implement their code into a basic Learner module I’m running into issues. I’m fairly certain it is their loss function, as it involves a complicated slew of metrics found within training if I’m not mistaken? Any help is appreciated to where I should get started here.