Fix assertion in GQA backward example to ensure correct tensor comparison for...

Fix assertion in GQA backward example to ensure correct tensor comparison for gradient validation (#568)

examples/flash_attention/example_gqa_bwd.py

@@ -342,7 +342,7 @@ def main(BATCH: int = 8,
 
     assert torch.allclose(O, O_ref, rtol=1e-2, atol=1e-2)
     torch.testing.assert_close(dV, dV_ref, rtol=1e-2, atol=1e-2)
-    # assert torch.allclose(dV, dV_ref, rtol=1e-2, atol=1e-2)
+    assert torch.allclose(dV, dV_ref, rtol=1e-2, atol=1e-2)
     assert torch.allclose(dK, dK_ref, rtol=1e-2, atol=1e-2)
     assert torch.allclose(dQ, dQ_ref, rtol=1e-2, atol=1e-2)