(self)
| 73 | |
| 74 | @skipUnless(has_einops, "Requires einops") |
| 75 | def test_flash_attention(self): |
| 76 | device = "cuda" if torch.cuda.is_available() else "cpu" |
| 77 | block = CABlock(spatial_dims=2, dim=64, num_heads=4, bias=True, flash_attention=True).to(device) |
| 78 | x = torch.randn(2, 64, 32, 32).to(device) |
| 79 | output = block(x) |
| 80 | self.assertEqual(output.shape, x.shape) |
| 81 | |
| 82 | @skipUnless(has_einops, "Requires einops") |
| 83 | def test_temperature_parameter(self): |