Fix remat failures in offload.py#13

copybara-service[bot] wants to merge 1 commit intomaingoogle/parallax:mainfrom

test_904235597google/parallax:test_904235597Copy head branch name to clipboard