def METHOD_NAME(self):
    grads = torch.randn_like(self.tst_inputs_q)
    time_mask_byte = torch.triu(
        torch.ones(
            self.tst_inputs_q.size(0), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8
        ),
        1,
    )
    time_mask_bool = time_mask_byte.to(torch.bool)
    ref_outputs, _ = self.ref_layer.forward(
        self.ref_inputs_q,
        self.ref_inputs_k,
        self.ref_inputs_k,
        key_padding_mask=None,
        need_weights=False,
        attn_mask=time_mask_bool,
        is_training=True,
    )
    tst_outputs, _ = self.tst_layer.forward(
        self.tst_inputs_q,
        self.tst_inputs_k,
        self.tst_inputs_k,
        key_padding_mask=None,
        need_weights=False,
        attn_mask=time_mask_byte,
        is_training=True,
    )
    self.ref_inputs_q.backward(grads)
    self.tst_inputs_q.backward(grads)
    torch.testing.assert_close(self.ref_inputs_q, self.tst_inputs_q, atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(self.ref_inputs_k, self.tst_inputs_k, atol=1e-5, rtol=1e-5)
    torch.testing.assert_close(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3)
    torch.testing.assert_close(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3)