fix: Coherent dtypes of updates with and without MultiSteps#1122

Merged

copybara-service[bot] merged 2 commits intogoogle-deepmind:mainfrom hlzl:main

Nov 5, 2024

+91-11

Commits on Oct 28, 2024

fix: updates with equivalent dtypes with and w/out MultiSteps.
hlzl
committed

Commits on Oct 29, 2024

fix: Add missing Loss.
hlzl
committed