fix qat_start bug in SparseML wrapper #17

fix qat_start bug in SparseML wrapper
neuralmagic · Aug 30, 2021 · 63b9a73 · 63b9a73
2 parents 372ffa9 + 056a9be
commit 63b9a73
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/sparse.py b/utils/sparse.py
@@ -119,7 +119,7 @@ def qat_active(self, epoch):
         if not self.enabled or not self.manager.quantization_modifiers:
             return False
 
-        qat_start = max([mod.start_epoch for mod in self.manager.quantization_modifiers])
+        qat_start = min([mod.start_epoch for mod in self.manager.quantization_modifiers])
 
         return qat_start < epoch + 1
 
@@ -136,4 +136,4 @@ def reset_best(self, epoch):
         qat_start = math.floor(max([mod.start_epoch for mod in self.manager.quantization_modifiers])) \
             if self.manager.quantization_modifiers else -1
 
-        return (pruning_start <= epoch <= pruning_end) or epoch == qat_start
+        return (pruning_start <= epoch <= pruning_end) or epoch == qat_start