refactor unet_2d tests (#13901)

Commit

56 days ago

refactor unet_2d tests (#13901) Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>