- am-training-step ctx-fwd ctx-bwd optimizer model loss-fn xprocedure
Performs one training step through the two contexts: the forward pass layer by layer, the loss, the backward pass, the optimizer step, and zeroing of the gradients. After the first step both contexts reuse the buffers they planned. Returns the loss as a lazy tensor.