Ahh, I was looking at the Block (which was what I found through google). Didn’t think it would have it’s own tensor wrapper as well. Thanks!
(Btw thanks for your image regression notebook. But I think I found an issue with using flip on the cat images here: Flipping and image regression with keypoints - #2 by Matsemann )