The model debugger was still crashing on preemption. This was due to the output of capture_intermediates returning pytrees with tuples. This when checkpointed would be turned into a dictionary with keys "0", "1", .... and would break upon reload. This CL fixes the issue by removing the unnecessary tuples from the pytrees.
The model debugger was still crashing on preemption. This was due to the output of capture_intermediates returning pytrees with tuples. This when checkpointed would be turned into a dictionary with keys "0", "1", .... and would break upon reload. This CL fixes the issue by removing the unnecessary tuples from the pytrees.

PiperOrigin-RevId: 441047399

google/init2winit