btw i created a pytorch rigl implementation some time ago, might be of use/inspiration: https://github.com/nollied/rigl-torch