mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-28 07:57:44 +00:00
fix: use trace count as index (#909)
This commit is contained in:
@@ -70,7 +70,7 @@ class LimitTimeCKPSelector(CheckpointSelector):
|
||||
return (-1,) # Continue with latest trial for new sub-trace
|
||||
|
||||
# Calculate elapsed time for current sub-trace
|
||||
elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count]
|
||||
elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count - 1]
|
||||
|
||||
if elapsed_time < self.time_limit_pre_trace:
|
||||
# Continue with current sub-trace
|
||||
@@ -81,7 +81,7 @@ class LimitTimeCKPSelector(CheckpointSelector):
|
||||
return (-1,)
|
||||
else:
|
||||
# Check if we've reached the maximum number of traces
|
||||
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
|
||||
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
|
||||
logger.info(
|
||||
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
|
||||
)
|
||||
@@ -89,11 +89,11 @@ class LimitTimeCKPSelector(CheckpointSelector):
|
||||
return (-1,)
|
||||
|
||||
# Time limit exceeded, start a new sub-trace
|
||||
self.sub_trace_start_times[trace.sub_trace_count + 1] = current_time
|
||||
self.sub_trace_start_times[trace.sub_trace_count] = current_time
|
||||
logger.info(
|
||||
f"Elapsed time {elapsed_time} exceeds time limit {self.time_limit_pre_trace}, jump to a new sub-trace"
|
||||
)
|
||||
logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
|
||||
logger.info(f"current sub-trace count: {trace.sub_trace_count}")
|
||||
return tuple() # Empty tuple signals starting a new sub-trace
|
||||
|
||||
|
||||
@@ -129,7 +129,7 @@ class SOTAJumpCKPSelector(CheckpointSelector):
|
||||
sota_count += 1
|
||||
if sota_count < self.SOTA_COUNT_THRESHOLD:
|
||||
# Check if we've reached the maximum number of traces
|
||||
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
|
||||
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
|
||||
logger.info(
|
||||
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
|
||||
)
|
||||
@@ -139,7 +139,7 @@ class SOTAJumpCKPSelector(CheckpointSelector):
|
||||
logger.info(
|
||||
f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump to a new sub-trace"
|
||||
)
|
||||
logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
|
||||
logger.info(f"current sub-trace count: {trace.sub_trace_count}")
|
||||
return ()
|
||||
else:
|
||||
logger.info(
|
||||
@@ -189,7 +189,7 @@ class BackJumpCKPSelector(CheckpointSelector):
|
||||
|
||||
if sota_count < self.SOTA_COUNT_THRESHOLD:
|
||||
# Check if we've reached the maximum number of traces before creating a new one
|
||||
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
|
||||
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
|
||||
logger.info(
|
||||
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
|
||||
)
|
||||
@@ -216,7 +216,7 @@ class BackJumpCKPSelector(CheckpointSelector):
|
||||
return (last_second_sota_idx,)
|
||||
else:
|
||||
# Check max trace limit again before creating a new trace
|
||||
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
|
||||
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
|
||||
logger.info(
|
||||
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
|
||||
)
|
||||
@@ -226,7 +226,7 @@ class BackJumpCKPSelector(CheckpointSelector):
|
||||
logger.info(
|
||||
f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump a new sub-trace"
|
||||
)
|
||||
logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
|
||||
logger.info(f"current sub-trace count: {trace.sub_trace_count}")
|
||||
return () # reboot a new sub-trace
|
||||
|
||||
else:
|
||||
|
||||
Reference in New Issue
Block a user