Comment by duchenne

Looks awesome. Can we get the same thing for pytorch?