From f8042757b34f579c391da48f701a26855d24bc50 Mon Sep 17 00:00:00 2001 From: Cassidy Laidlaw Date: Fri, 18 Nov 2022 12:01:19 -0800 Subject: [PATCH] Fix formatting Signed-off-by: Cassidy Laidlaw --- rllib/policy/sample_batch.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rllib/policy/sample_batch.py b/rllib/policy/sample_batch.py index a1c04cca619c..81c9bf16c534 100644 --- a/rllib/policy/sample_batch.py +++ b/rllib/policy/sample_batch.py @@ -56,10 +56,10 @@ def attempt_count_timesteps(tensor_dict: dict): and not (tf and tf.is_tensor(tensor_dict[SampleBatch.SEQ_LENS])) and len(tensor_dict[SampleBatch.SEQ_LENS]) > 0 ): - if torch and torch.is_tensor(tensor_dict[SampleBatch.SEQ_LENS]): - return tensor_dict[SampleBatch.SEQ_LENS].sum().item() - else: - return sum(tensor_dict[SampleBatch.SEQ_LENS]) + if torch and torch.is_tensor(tensor_dict[SampleBatch.SEQ_LENS]): + return tensor_dict[SampleBatch.SEQ_LENS].sum().item() + else: + return sum(tensor_dict[SampleBatch.SEQ_LENS]) for k, v in copy_.items(): assert isinstance(k, str), tensor_dict