Fix ModelParallel OOM issue during weight loading#21723

Open

amitsrivastava78 wants to merge 10 commits intokeras-team:masterfrom

amitsrivastava78:pytorch-fix-rebased

Commits on Oct 7, 2025

Fix ModelParallel OOM issue during weight loading
amitsrivastava78
committed
Fix PyTorch backend tensor conversion and refactor variable loading
amitsrivastava78
committed
Optimize JAX backend variable initialization and memory management
amitsrivastava78
committed

Commits on Oct 9, 2025

Update all layer load_own_variables methods to use _direct_assign
amitsrivastava78
committed
Add shape validation to layer load_own_variables methods for consistent error handling
amitsrivastava78
committed