- cast gradients and inputs to weight.dtype instead of hardcoded bf16 - single code path covers bf16 and fp32 models, no branch needed - gradient dtype now matches the leaf parameter dtype exactly
- cast gradients and inputs to weight.dtype instead of hardcoded bf16 - single code path covers bf16 and fp32 models, no branch needed - gradient dtype now matches the leaf parameter dtype exactly