Fix remat failures in offload.py#12

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

test_895541719google/parallax:test_895541719Copy head branch name to clipboard