Demo equivalence - tmp

[gordicaleksa]

cc: @karpathy

here is the minimal change on top of my PR that gives equivalent code to nano llama 3 reference.py from commit karpathy/nano-llama31@d0dfb06

Steps:

1. Check out this PR
2. Check out the commit I linked above from nano llama 3 (the only thing you've introduced since is passing in a generator, so functionality-wise this is an equivalent commit)
3. Comment out nn.Linear & nn.Embedding's reset params function in PyTorch constructor
4. Run and see same samples