[SPMD] Fix TPU CI after #5478

Commit

2 years ago

[SPMD] Fix TPU CI after #5478 Summary: Let's fix all TPU CI failures after #5478. Test Plan: TPU CI