An implementation of the "A simple neural network module for relational reasoning" paper in NIPS 2017