Since the backbone of the implementation is packed away in the library import, I felt it didn't quite show the code & variable interaction well enough.

Don't get me wrong. It is useful & concise like you mentioned. But your target audience is beginners & adopters, and it makes it no different from another framework such as Fastai (I have major gripes with them. It has a much bottled-in experience)

To be true to walkthroughs, please consider designing helper functions rather than using your framework. Admittedly it may not be as beautiful, but eventually your users will be more appreciative of the extra mile you go into making things transparent & similar to PyTorch docs.