def METHOD_NAME(self):
    expected_x = torch.tensor(
        [
            [
                [
                    [-0.7193, -0.4033, -0.5966],
                    [-0.8567, 1.1006, -1.0712],
                    [-0.5663, 0.3731, -0.8920],
                ]
            ]
        ]
    )
    x = self.MHA.rel_shift(self.sample_x)
    self.assertTrue(
        np.allclose(
            expected_x.cpu().detach().numpy(),
            x.cpu().detach().numpy(),
            atol=1e-4,
        )
    )