The architectures library is working right now because there's only one architecture. This will actually break for more than one architecture, refer to the implementation in BBHNet for how to fix this. Better yet, just create a dedicated ML4GW/architectures repo that contains architectures for various purposes and this wrapper.