Revert cleanup approach, restore forward-pass fix for embed_tokens
The GPU cannot fit Gemma4-E2B entirely with device_map="auto", so
accelerate offloads embed_tokens to CPU on every model load — not just
under accumulated memory pressure. The cleanup/tearDown approach
doesn't help; it just forces model reloads that hit the same issue.
The correct fix is to use embed_tokens.forward() instead of direct
.weight access: the forward pass triggers accelerate's AlignDevicesHook
which transiently loads weights from CPU to GPU before the computation.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>