Move `disable_v2_behaviour()` to `setUp` to enable running tests in test suites.

PiperOrigin-RevId: 454823596
Change-Id: I893dc051a5f21d53905acee45a4dcb5c5609daf9

alphafold/model/tf/protein_features_test.py

@@ -27,6 +27,10 @@ def _random_bytes():
 
 class FeaturesTest(parameterized.TestCase, tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    tf.disable_v2_behavior()
+
   def testFeatureNames(self):
     self.assertEqual(len(protein_features.FEATURE_SIZES),
                      len(protein_features.FEATURE_TYPES))
@@ -47,5 +51,4 @@ class FeaturesTest(parameterized.TestCase, tf.test.TestCase):
 
 
 if __name__ == '__main__':
-  tf.disable_v2_behavior()
   absltest.main()

alphafold/model/tf/shape_helpers_test.py

@@ -21,6 +21,10 @@ import tensorflow.compat.v1 as tf
 
 class ShapeTest(tf.test.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    tf.disable_v2_behavior()
+
   def test_shape_list(self):
     """Test that shape_list can allow for reshaping to dynamic shapes."""
     a = tf.zeros([10, 4, 4, 2])
@@ -35,5 +39,4 @@ class ShapeTest(tf.test.TestCase):
 
 
 if __name__ == '__main__':
-  tf.disable_v2_behavior()
   tf.test.main()