Address comments

apache · Feb 14, 2020 · 1acf051 · 1acf051
1 parent e2f417e
commit 1acf051
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/tvm/autotvm/tuner/index_based_tuner.py b/python/tvm/autotvm/tuner/index_based_tuner.py
@@ -60,7 +60,7 @@ def next_batch(self, batch_size):
         ret = []
         for _ in range(batch_size):
             if self.counter >= self.range_length:
-                continue
+                break
             index = self.counter + self.index_offset
             ret.append(self.task.config_space.get(index))
             self.counter = self.counter + 1
@@ -115,7 +115,8 @@ def next_batch(self, batch_size):
         return ret
 
     def __getstate__(self):
-        return {"visited": self.counter}
+        return {"visited": self.visited}
 
     def __setstate__(self, state):
-        self.counter = state['visited']
+        self.visited = state['visited']
+        self.counter = len(self.visited)