This repo is still wip and a cleaned up brach will be created to reproduce our new paper https://openreview.net/forum?id=0ku2hIm4BS Overall the code to reproduce it is in the generalization folder.