Revert code changes to mutltihead_attn tests

apex/contrib/test/multihead_attn/test_encdec_multihead_attn.py

@@ -47,6 +47,8 @@ class EncdecMultiheadAttnTest(unittest.TestCase):
                                         dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
 
     def test_encdec_multihead_attn(self) :
+        grads         = torch.randn_like(self.tst_inputs_q)
+
         ref_outputs,_ = self.ref_layer.forward(self.ref_inputs_q,
                                                self.ref_inputs_k,
                                                self.ref_inputs_k,
@@ -62,15 +64,13 @@ class EncdecMultiheadAttnTest(unittest.TestCase):
                                                need_weights=False,
                                                attn_mask=None,
                                                is_training=True)
+
+        self.ref_inputs_q.backward(grads)
+        self.tst_inputs_q.backward(grads)
+
         self.assertTrue(torch.allclose(self.ref_inputs_q,  self.tst_inputs_q,  atol=1e-5, rtol=1e-5))
         self.assertTrue(torch.allclose(self.ref_inputs_k,  self.tst_inputs_k,  atol=1e-5, rtol=1e-5))
         self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
-
-        with torch.no_grad():
-            ref_grads         = torch.randn_like(ref_outputs)
-            tst_grads = ref_grads.clone()
-        ref_outputs.backward(ref_grads)
-        tst_outputs.backward(tst_grads)
         self.assertTrue(torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3))
     
     def test_encdec_multihead_attn_time_mask(self) :

apex/contrib/test/multihead_attn/test_self_multihead_attn.py

@@ -42,7 +42,9 @@ class SelfMultiheadAttnTest(unittest.TestCase):
         self.tst_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim,
                                       dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
 
-    def test_self_multihead_attn(self) :
+    def test_self_multihead_attn(self):
+        grads         = torch.randn_like(self.tst_inputs)
+
         ref_outputs,_ = self.ref_layer.forward(self.ref_inputs,
                                                self.ref_inputs,
                                                self.ref_inputs,
@@ -59,15 +61,11 @@ class SelfMultiheadAttnTest(unittest.TestCase):
                                                attn_mask=None,
                                                is_training=True)
 
+        self.ref_inputs.backward(grads)
+        self.tst_inputs.backward(grads)
+
         self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
         self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
-
-        with torch.no_grad():
-            ref_grads         = torch.randn_like(self.tst_inputs)
-            tst_grads         = ref_grads.clone()
-
-        ref_outputs.backward(ref_grads)
-        tst_outputs.backward(tst_grads)
         self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
 
     def test_self_multihead_attn_time_mask(self) :