fix qat_start bug in SparseML wrapper

neuralmagic · Apr 14, 2022 · 4c24632 · 4c24632
1 parent 8255139
commit 4c24632
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