This code here is strongly inspired by the implementation of FermiNet (Copyright 2020 DeepMind Technologies Limited.)