diff --git a/gematria/model/python/BUILD.bazel b/gematria/model/python/BUILD.bazel
index bbf44300..de18cba4 100644
--- a/gematria/model/python/BUILD.bazel
+++ b/gematria/model/python/BUILD.bazel
@@ -50,9 +50,6 @@ gematria_py_test(
     size = "small",
     timeout = "moderate",
     srcs = ["loss_utils_test.py"],
-    tags = [
-        "manual",
-    ],
     deps = [
         ":loss_utils",
         ":options",
diff --git a/gematria/model/python/loss_utils_test.py b/gematria/model/python/loss_utils_test.py
index 6f8bd5a1..462611c7 100644
--- a/gematria/model/python/loss_utils_test.py
+++ b/gematria/model/python/loss_utils_test.py
@@ -15,7 +15,7 @@
 from gematria.model.python import loss_utils
 from gematria.model.python import options
 import numpy as np
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 
 class LossComputationTest(tf.test.TestCase):
@@ -62,24 +62,18 @@ def test_unscaled_loss(self):
         dtype=self.dtype,
     )
 
-    mse_tensor = loss.mean_squared_error
-    mae_tensor = loss.mean_absolute_error
-    huber_tensor = loss.loss_tensor(
+    mse = loss.mean_squared_error
+    mae = loss.mean_absolute_error
+    huber = loss.loss_tensor(
         options.ErrorNormalization.NONE, options.LossType.HUBER
     )
-    percentiles_tensor = loss.absolute_error_percentiles
-    with self.session() as sess:
-      mse, mae, huber, percentiles = sess.run(
-          (mse_tensor, mae_tensor, huber_tensor, percentiles_tensor)
-      )
-      self.assertNear(
-          float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6
-      )
-      self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
-      self.assertNear(
-          float(huber), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
-      )
-      self.assertAllEqual(percentiles, ((0,), (0,), (0.5,), (3,), (4,)))
+    percentiles = loss.absolute_error_percentiles
+    self.assertNear(float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6)
+    self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
+    self.assertNear(
+        float(huber), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
+    )
+    self.assertAllEqual(percentiles, ((0,), (0,), (0.5,), (3,), (4,)))
 
   def test_percentage_loss(self):
     loss = loss_utils.LossComputation(
@@ -90,18 +84,14 @@ def test_percentage_loss(self):
         dtype=self.dtype,
     )
 
-    mspe_tensor = loss.mean_squared_percentage_error
-    mape_tensor = loss.mean_absolute_percentage_error
-    percentiles_tensor = loss.absolute_percentage_error_percentiles
-    with self.session() as sess:
-      mspe, mape, percentiles = sess.run(
-          (mspe_tensor, mape_tensor, percentiles_tensor)
-      )
-      self.assertAlmostEqual(
-          float(mspe), ((3 / 4) ** 2 + 0 + 0 + 4**2 + (0.5 / 2) ** 2) / 5
-      )
-      self.assertAlmostEqual(float(mape), (3 / 4 + 0 + 0 + 4 + 0.5 / 2) / 5)
-      self.assertAllEqual(percentiles, ((0,), (0.5 / 2,), (0.5 / 2,), (4,)))
+    mspe = loss.mean_squared_percentage_error
+    mape = loss.mean_absolute_percentage_error
+    percentiles = loss.absolute_percentage_error_percentiles
+    self.assertAlmostEqual(
+        float(mspe), ((3 / 4) ** 2 + 0 + 0 + 4**2 + (0.5 / 2) ** 2) / 5
+    )
+    self.assertAlmostEqual(float(mape), (3 / 4 + 0 + 0 + 4 + 0.5 / 2) / 5)
+    self.assertAllEqual(percentiles, ((0,), (0.5 / 2,), (0.5 / 2,), (4,)))
 
   def test_normalized_loss_when_expected_value_greater_than_one(self):
     actual_outputs = tf.constant(
@@ -115,25 +105,22 @@ def test_normalized_loss_when_expected_value_greater_than_one(self):
         actual_outputs, expected_outputs, mask, dtype=self.dtype
     )
 
-    mean_absolute_error_tensor = loss.loss_tensor(
+    mean_absolute_error = loss.loss_tensor(
         options.ErrorNormalization.EXPECTED_VALUE_GREATER_THAN_ONE,
         options.LossType.MEAN_ABSOLUTE_ERROR,
     )
-    mean_squared_error_tensor = loss.loss_tensor(
+    mean_squared_error = loss.loss_tensor(
         options.ErrorNormalization.EXPECTED_VALUE_GREATER_THAN_ONE,
         options.LossType.MEAN_SQUARED_ERROR,
     )
-    with self.session() as sess:
-      mean_absolute_error = sess.run(mean_absolute_error_tensor)
-      mean_squared_error = sess.run(mean_squared_error_tensor)
-      self.assertAlmostEqual(
-          float(mean_absolute_error), (0.3 + 1.5 + 0.0 + 0.5 + 2.0) / 5
-      )
-      self.assertAlmostEqual(
-          float(mean_squared_error),
-          (0.3**2 + 1.5**2 + 0.0 + 0.5**2 + 2.0**2) / 5,
-          delta=1e-6,
-      )
+    self.assertAlmostEqual(
+        float(mean_absolute_error), (0.3 + 1.5 + 0.0 + 0.5 + 2.0) / 5
+    )
+    self.assertAlmostEqual(
+        float(mean_squared_error),
+        (0.3**2 + 1.5**2 + 0.0 + 0.5**2 + 2.0**2) / 5,
+        delta=1e-6,
+    )
 
   def test_with_no_percentiles(self):
     loss = loss_utils.LossComputation(
@@ -142,14 +129,10 @@ def test_with_no_percentiles(self):
         self.full_mask,
         dtype=self.dtype,
     )
-    abs_percentiles_tensor = loss.absolute_error_percentiles
-    pct_percentiles_tenros = loss.absolute_percentage_error_percentiles
-    with self.session() as sess:
-      abs_percentiles, pct_percentiles = sess.run(
-          (abs_percentiles_tensor, pct_percentiles_tenros)
-      )
-      self.assertAllEqual(abs_percentiles, [])
-      self.assertAllEqual(pct_percentiles, [])
+    abs_percentiles = loss.absolute_error_percentiles
+    pct_percentiles = loss.absolute_percentage_error_percentiles
+    self.assertAllEqual(abs_percentiles, [])
+    self.assertAllEqual(pct_percentiles, [])
 
   def test_multi_task(self):
     self.assertEqual(self.multitask_actual_outputs.shape, (5, 2))
@@ -163,38 +146,34 @@ def test_multi_task(self):
         dtype=self.dtype,
     )
 
-    mse_tensor = loss.mean_squared_error
-    mae_tensor = loss.mean_absolute_error
-    huber_tensor = loss.loss_tensor(
+    mse = loss.mean_squared_error
+    mae = loss.mean_absolute_error
+    huber = loss.loss_tensor(
         options.ErrorNormalization.NONE, options.LossType.HUBER
     )
-    percentiles_tensor = loss.absolute_error_percentiles
-    with self.session() as sess:
-      mse, mae, huber, percentiles = sess.run(
-          (mse_tensor, mae_tensor, huber_tensor, percentiles_tensor)
-      )
-      self.assertAllClose(
-          mse,
-          (
-              (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
-              (6**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
-          ),
-          1e-6,
-      )
-      self.assertAllClose(
-          mae, ((3 + 0 + 0 + 4 + 0.5) / 5, (6 + 0 + 0 + 4 + 0.5) / 5), 1e-6
-      )
-      self.assertAllClose(
-          huber,
-          (
-              (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
-              (5.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
-          ),
-          1e-6,
-      )
-      self.assertAllEqual(
-          percentiles, ((0, 0), (0, 0), (0.5, 0.5), (3, 4), (4, 6))
-      )
+    percentiles = loss.absolute_error_percentiles
+    self.assertAllClose(
+        mse,
+        (
+            (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
+            (6**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
+        ),
+        1e-6,
+    )
+    self.assertAllClose(
+        mae, ((3 + 0 + 0 + 4 + 0.5) / 5, (6 + 0 + 0 + 4 + 0.5) / 5), 1e-6
+    )
+    self.assertAllClose(
+        huber,
+        (
+            (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+            (5.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+        ),
+        1e-6,
+    )
+    self.assertAllEqual(
+        percentiles, ((0, 0), (0, 0), (0.5, 0.5), (3, 4), (4, 6))
+    )
 
   def test_multi_task_with_mask(self):
     loss = loss_utils.LossComputation(
@@ -216,173 +195,133 @@ def test_multi_task_with_mask(self):
         dtype=self.dtype,
     )
 
-    mse_tensor = loss.mean_squared_error
-    mspe_tensor = loss.mean_squared_percentage_error
-    mae_tensor = loss.mean_absolute_error
-    mape_tensor = loss.mean_absolute_percentage_error
-    huber_tensor = loss.loss_tensor(
+    mse = loss.mean_squared_error
+    mspe = loss.mean_squared_percentage_error
+    mae = loss.mean_absolute_error
+    mape = loss.mean_absolute_percentage_error
+    huber = loss.loss_tensor(
         options.ErrorNormalization.NONE, options.LossType.HUBER
     )
-    huber_percentage_tensor = loss.loss_tensor(
+    huber_percentage = loss.loss_tensor(
         options.ErrorNormalization.PERCENTAGE_ERROR, options.LossType.HUBER
     )
-    percentile_tensor = loss.absolute_error_percentiles
-    with self.session() as sess:
-      mse, mspe, mae, mape, huber, huber_percentage, percentiles = sess.run((
-          mse_tensor,
-          mspe_tensor,
-          mae_tensor,
-          mape_tensor,
-          huber_tensor,
-          huber_percentage_tensor,
-          percentile_tensor,
-      ))
-      self.assertAllClose(
-          mse,
-          ((3**2 + 27**2 + 0.5**2) / 3, (6**2 + 0.1**2) / 2),
-          1e-6,
-      )
-      self.assertAllClose(
-          mspe,
-          (
-              ((3 / 4) ** 2 + (27 / 30) ** 2 + (0.5 / 2) ** 2) / 3,
-              ((6 / 14) ** 2 + (0.1 / 12) ** 2) / 2,
-          ),
-          1e-6,
-      )
-      self.assertAllClose(mae, ((3 + 27 + 0.5) / 3, (6 + 0.1) / 2), 1e-6)
-      self.assertAllClose(
-          mape, ((3 / 4 + 27 / 30 + 0.5 / 2) / 3, (6 / 14 + 0.1 / 12) / 2), 1e-6
-      )
-      self.assertAllClose(
-          huber,
-          ((2.5 + 26.5 + (0.5**2) / 2) / 3, (5.5 + (0.1**2) / 2) / 2),
-          1e-6,
-      )
-      self.assertAllClose(
-          huber_percentage,
-          (
-              (((3 / 4) ** 2) / 2 + ((27 / 30) ** 2) / 2 + ((0.5 / 2) ** 2) / 2)
-              / 3,
-              (((6 / 14) ** 2) / 2 + ((0.1 / 12) ** 2) / 2) / 2,
-          ),
-          1e-6,
-      )
-      self.assertAllClose(
-          percentiles,
-          ((0.5, 0.1), (0.5, 0.1), (3, 0.1), (27, 6), (27, 6)),
-          1e-6,
-      )
+    percentiles = loss.absolute_error_percentiles
+    self.assertAllClose(
+        mse,
+        ((3**2 + 27**2 + 0.5**2) / 3, (6**2 + 0.1**2) / 2),
+        1e-6,
+    )
+    self.assertAllClose(
+        mspe,
+        (
+            ((3 / 4) ** 2 + (27 / 30) ** 2 + (0.5 / 2) ** 2) / 3,
+            ((6 / 14) ** 2 + (0.1 / 12) ** 2) / 2,
+        ),
+        1e-6,
+    )
+    self.assertAllClose(mae, ((3 + 27 + 0.5) / 3, (6 + 0.1) / 2), 1e-6)
+    self.assertAllClose(
+        mape, ((3 / 4 + 27 / 30 + 0.5 / 2) / 3, (6 / 14 + 0.1 / 12) / 2), 1e-6
+    )
+    self.assertAllClose(
+        huber,
+        ((2.5 + 26.5 + (0.5**2) / 2) / 3, (5.5 + (0.1**2) / 2) / 2),
+        1e-6,
+    )
+    self.assertAllClose(
+        huber_percentage,
+        (
+            (((3 / 4) ** 2) / 2 + ((27 / 30) ** 2) / 2 + ((0.5 / 2) ** 2) / 2)
+            / 3,
+            (((6 / 14) ** 2) / 2 + ((0.1 / 12) ** 2) / 2) / 2,
+        ),
+        1e-6,
+    )
+    self.assertAllClose(
+        percentiles,
+        ((0.5, 0.1), (0.5, 0.1), (3, 0.1), (27, 6), (27, 6)),
+        1e-6,
+    )
 
   def test_unknown_shape(self):
-    actual_output = tf.placeholder(self.dtype, shape=(None, 1))
-    expected_output = tf.placeholder(self.dtype, shape=(None, 1))
-    mask = tf.placeholder(tf.dtypes.bool, shape=(None, 1))
     percentile_ranks = (0, 50, 75, 100)
 
     loss = loss_utils.LossComputation(
-        actual_output,
-        expected_output,
-        mask,
+        self.actual_outputs,
+        self.expected_outputs,
+        tf.ones_like(self.actual_outputs, dtype=bool),
         percentile_ranks=percentile_ranks,
         dtype=self.dtype,
     )
 
-    mse_tensor = loss.mean_squared_error
-    mae_tensor = loss.mean_absolute_error
-    huber_tensor = loss.loss_tensor(
+    mse = loss.mean_squared_error
+    mae = loss.mean_absolute_error
+    huber = loss.loss_tensor(
         options.ErrorNormalization.NONE, options.LossType.HUBER
     )
-    percentiles_tensor = loss.absolute_error_percentiles
-
-    self.assertEqual(mse_tensor.shape, (1,))
-    self.assertEqual(mae_tensor.shape, (1,))
-    self.assertEqual(huber_tensor.shape, (1,))
-    self.assertEqual(percentiles_tensor.shape, (len(percentile_ranks), 1))
-
-    feed_dict = {
-        actual_output: self.actual_outputs_array,
-        expected_output: self.expected_outputs_array,
-        mask: np.ones_like(self.actual_outputs_array, dtype=bool),
-    }
-    with self.session() as sess:
-      mse, mae, huber, percentiles = sess.run(
-          (mse_tensor, mae_tensor, huber_tensor, percentiles_tensor),
-          feed_dict=feed_dict,
-      )
-      self.assertNear(
-          float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6
-      )
-      self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
-      self.assertNear(
-          float(huber), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
-      )
-      self.assertAllEqual(percentiles, ((0,), (0.5,), (3,), (4,)))
+    percentiles = loss.absolute_error_percentiles
+
+    self.assertEqual(mse.shape, (1,))
+    self.assertEqual(mae.shape, (1,))
+    self.assertEqual(huber.shape, (1,))
+    self.assertEqual(percentiles.shape, (len(percentile_ranks), 1))
+
+    self.assertNear(float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6)
+    self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
+    self.assertNear(
+        float(huber), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
+    )
+    self.assertAllEqual(percentiles, ((0,), (0.5,), (3,), (4,)))
 
   def test_multi_task_unknown_shape(self):
     num_tasks = 2
-    actual_output = tf.placeholder(self.dtype, shape=(None, num_tasks))
-    expected_output = tf.placeholder(self.dtype, shape=(None, num_tasks))
-    mask = tf.placeholder(tf.dtypes.bool, shape=(None, num_tasks))
     percentile_ranks = (0, 50, 75, 100)
 
     loss = loss_utils.LossComputation(
-        actual_output,
-        expected_output,
-        mask,
+        self.multitask_actual_outputs,
+        self.multitask_expected_outputs,
+        tf.ones_like(self.multitask_actual_outputs_array, dtype=bool),
         percentile_ranks=percentile_ranks,
         dtype=self.dtype,
     )
 
-    mse_tensor = loss.mean_squared_error
-    mae_tensor = loss.mean_absolute_error
-    huber_tensor = loss.loss_tensor(
+    mse = loss.mean_squared_error
+    mae = loss.mean_absolute_error
+    huber = loss.loss_tensor(
         options.ErrorNormalization.NONE, options.LossType.HUBER
     )
-    percentiles_tensor = loss.absolute_error_percentiles
-
-    self.assertEqual(mse_tensor.shape, (num_tasks,))
-    self.assertEqual(mae_tensor.shape, (num_tasks,))
-    self.assertEqual(huber_tensor.shape, (num_tasks,))
-    self.assertEqual(
-        percentiles_tensor.shape, (len(percentile_ranks), num_tasks)
-    )
-
-    feed_dict = {
-        actual_output: self.multitask_actual_outputs_array,
-        expected_output: self.multitask_expected_outputs_array,
-        mask: np.ones_like(self.multitask_actual_outputs_array, dtype=bool),
-    }
-    with self.session() as sess:
-      mse, mae, huber, percentiles = sess.run(
-          (mse_tensor, mae_tensor, huber_tensor, percentiles_tensor),
-          feed_dict=feed_dict,
-      )
-      self.assertNDArrayNear(
-          mse,
-          np.array((
-              (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
-              (6**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
-          )),
-          1e-6,
-      )
-      self.assertNDArrayNear(
-          mae, (((3 + 0 + 0 + 4 + 0.5) / 5, (6 + 0 + 0 + 4 + 0.5) / 5)), 1e-6
-      )
-      self.assertNDArrayNear(
-          huber,
-          ((
-              (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
-              (5.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
-          )),
-          1e-6,
-      )
-      self.assertAllEqual(percentiles, ((0, 0), (0.5, 0.5), (3, 4), (4, 6)))
+    percentiles = loss.absolute_error_percentiles
+
+    self.assertEqual(mse.shape, (num_tasks,))
+    self.assertEqual(mae.shape, (num_tasks,))
+    self.assertEqual(huber.shape, (num_tasks,))
+    self.assertEqual(percentiles.shape, (len(percentile_ranks), num_tasks))
+
+    self.assertNDArrayNear(
+        mse,
+        np.array((
+            (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
+            (6**2 + 0 + 0 + 4**2 + 0.5**2) / 5,
+        )),
+        1e-6,
+    )
+    self.assertNDArrayNear(
+        mae, (((3 + 0 + 0 + 4 + 0.5) / 5, (6 + 0 + 0 + 4 + 0.5) / 5)), 1e-6
+    )
+    self.assertNDArrayNear(
+        huber,
+        ((
+            (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+            (5.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+        )),
+        1e-6,
+    )
+    self.assertAllEqual(percentiles, ((0, 0), (0.5, 0.5), (3, 4), (4, 6)))
 
   def test_single_task_unknown_shape(self):
     num_tasks = 1
-    actual_output = tf.placeholder(self.dtype, shape=(None, num_tasks))
-    expected_output = tf.placeholder(self.dtype, shape=(None, num_tasks))
+    actual_output = tf.reshape(self.actual_outputs_array, (-1, 1))
+    expected_output = tf.reshape(self.expected_outputs_array, (-1, 1))
     mask = tf.ones_like(actual_output, tf.dtypes.bool)
     percentile_ranks = (0, 50, 75, 100)
 
@@ -394,39 +333,25 @@ def test_single_task_unknown_shape(self):
         dtype=self.dtype,
     )
 
-    mse_tensor = loss.mean_squared_error
-    mae_tensor = loss.mean_absolute_error
-    huber_tensor = loss.loss_tensor(
+    mse = loss.mean_squared_error
+    mae = loss.mean_absolute_error
+    huber = loss.loss_tensor(
         options.ErrorNormalization.NONE, options.LossType.HUBER
     )
-    percentiles_tensor = loss.absolute_error_percentiles
-
-    self.assertEqual(mse_tensor.shape, (num_tasks,))
-    self.assertEqual(mae_tensor.shape, (num_tasks,))
-    self.assertEqual(huber_tensor.shape, (num_tasks,))
-    self.assertEqual(
-        percentiles_tensor.shape, (len(percentile_ranks), num_tasks)
-    )
-
-    feed_dict = {
-        actual_output: self.actual_outputs_array.reshape((-1, 1)),
-        expected_output: self.expected_outputs_array.reshape((-1, 1)),
-    }
-    with self.session() as sess:
-      mse, mae, huber, percentiles = sess.run(
-          (mse_tensor, mae_tensor, huber_tensor, percentiles_tensor),
-          feed_dict=feed_dict,
-      )
-      self.assertNear(
-          float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6
-      )
-      self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
-      self.assertNear(
-          float(huber), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
-      )
-      self.assertAllEqual(percentiles, ((0,), (0.5,), (3,), (4,)))
+    percentiles = loss.absolute_error_percentiles
+
+    self.assertEqual(mse.shape, (num_tasks,))
+    self.assertEqual(mae.shape, (num_tasks,))
+    self.assertEqual(huber.shape, (num_tasks,))
+    self.assertEqual(percentiles.shape, (len(percentile_ranks), num_tasks))
+
+    self.assertNear(float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6)
+    self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
+    self.assertNear(
+        float(huber), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
+    )
+    self.assertAllEqual(percentiles, ((0,), (0.5,), (3,), (4,)))
 
 
 if __name__ == '__main__':
-  tf.disable_v2_behavior()
   tf.test.main()