Comment by lucidrains

gradtree is cool! thanks for the share!