Learn · 4

Training

Zero the gradients, forward, loss, backward, step. The one thing torch does not have is scope(), and without it the device fills up in seconds.

Why scope exists

A single training step creates thousands of intermediate buffers. JavaScript's garbage collector does not release GPU memory promptly — it does not know that a small object owns a large allocation. scope() frees everything created inside it when it exits; keepAlive() marks what must survive. TF.js calls the same idea tidy.

Run this block and watch the numbers, then delete the scope wrapper and run it again.

Note that zeroGrad() is the first line inside the scope, not an afterthought. Gradients from the previous step were allocated in the previous scope; accumulating into them after that scope closed reads a buffer that has already been handed to something else. borch throws there rather than returning whatever now sits in that memory.

The loop

Data with a known answer is the honest way to see whether training works: if the loop cannot recover y = 3x + 2, nothing else it reports means anything.

This block carries the same loop twice. Run javascript, then press python and run it again: the losses agree to the last digit, because the two surfaces dispatch to the same WGSL. That is the claim this site makes about the Python route, and it is cheaper to check than to believe.

Optimizers and schedules

SGD, Adam, RMSprop and a dozen more are here, along with the learning-rate schedulers. Parameter groups work the way torch's do — per-group learning rates, and weight decay you can leave off the biases.