Add tests for 2D UNet blocks (#1945)

Commit

2 years ago

Add tests for 2D UNet blocks (#1945) * test unet blocks 2d * change to randn_tensor * mps flaky