Function train_step_adamw 

pub fn train_step_adamw(
    graph: &mut Graph,
    optimizer: &mut AdamW,
    prediction: NodeId,
    target: NodeId,
    trainable_nodes: &[NodeId],
) -> Result<f32, ModelError>

Runs one full train step: loss forward, backward, and AdamW updates.