unit test

engines/python/setup/djl_python/scheduler/seq_batch_scheduler.py

@@ -43,8 +43,7 @@ def __init__(self, lm_block: LMBlock, default_search_algorithm: str,
         self.results: Dict[int, List[int]] = defaultdict(list)
 
         self.seq_batchers: Dict[
-            Type[SeqBatcher]:List[SeqBatcher]] = defaultdict(
-                list)  # {key: List[SeqBatcher]}
+            Type[SeqBatcher]:List[SeqBatcher]] = defaultdict(list)
 
         self.lru_kv_cache = OrderedDict()
         self.lru_max_size = 10

engines/python/setup/djl_python/tests/test_scheduler.py

@@ -161,10 +161,6 @@ def test_sampling_scheduler(self):
 
         results = scheduler.collect_results()
 
-        print('=====++++++++++==============')
-        for i, ret in results.items():
-            print('\n{}:'.format(i), tokenizer.decode(ret[:30]))
-
         assert tokenizer.decode(results[0][:30]) == "Memories follow me left and right. I can't remember the last " \
                                                     "time I had to do this, and I'm glad that I did.\n"
         assert tokenizer.decode(
@@ -174,8 +170,6 @@ def test_sampling_scheduler(self):
             results[2][:30]
         ) == "There's a time that I remember, when I did not know what to do with my life. I was in a very bad mood. I was"
 
-        print('======+++++++++++++=========')
-
     def test_contrastive_scheduler(self):
         model_id = "gpt2"
         model = GPT2LMHeadModel.from_pretrained(model_id)