fix linting

Author: priyakasimbeg

tests/test_jax_utils.py

@@ -215,7 +215,6 @@ def test_jitted_updates(self, dropout_rate, mode):
           x,
           train=train,
           rngs={"dropout": dropout_rng})
-      return y1
 
     for d in [i * 0.1 * dropout_rate for i in range(0, 11)]:
       y2 = jitted_custom_apply(
@@ -225,8 +224,6 @@ def test_jitted_updates(self, dropout_rate, mode):
           dropout_rate=d,
           rngs={"dropout": dropout_rng},
       )
-    return y2
-
     assert jnp.allclose(y1, y2, atol=1e-3, rtol=1e-3)