Add a convergence test for training UNet3D on TPU v3-8

[thisisalbertliang]

@thisisalbertliang developed a UNet3D PyTorch/XLA codebase that can run on both TPU and GPU.
To enhance our TPU test suite, we hope to add a convergence test for this UNet3D model on TPU v3-8