yapf . -ir

google · Sep 6, 2024 · f5b6b6f · f5b6b6f
1 parent bf12cee
commit f5b6b6f
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/compiler_opt/tools/combine_tfa_policies_lib_test.py b/compiler_opt/tools/combine_tfa_policies_lib_test.py
@@ -30,9 +30,7 @@ class AddOnePolicy(tf_agents.policies.TFPolicy):
   """Test policy which adds one to obs feature."""
 
   def __init__(self):
-    obs_spec = {
-        'obs': tensor_spec.TensorSpec(shape=(1,), dtype=tf.int64)
-    }
+    obs_spec = {'obs': tensor_spec.TensorSpec(shape=(1,), dtype=tf.int64)}
     time_step_spec = ts.time_step_spec(obs_spec)
 
     act_spec = tensor_spec.TensorSpec(shape=(1,), dtype=tf.int64)
@@ -56,9 +54,7 @@ class SubtractOnePolicy(tf_agents.policies.TFPolicy):
   """Test policy which subtracts one to obs feature."""
 
   def __init__(self):
-    obs_spec = {
-        'obs': tensor_spec.TensorSpec(shape=(1,), dtype=tf.int64)
-    }
+    obs_spec = {'obs': tensor_spec.TensorSpec(shape=(1,), dtype=tf.int64)}
     time_step_spec = ts.time_step_spec(obs_spec)
 
     act_spec = tensor_spec.TensorSpec(shape=(1,), dtype=tf.int64)