Modified 'post_round' function condition (#23)

Author: nastiapetrovych

src/pquant/core/keras/layers.py

@@ -1347,11 +1347,12 @@ def call(self, x, training=None):
 
 
 def call_post_round_functions(model, rewind, rounds, r):
+    last_round = (r == rounds - 1)
     if rewind == "round":
         rewind_weights_functions(model)
-    elif rewind == "post-ticket-search" and r == rounds - 1:
+    elif rewind == "post-ticket-search" and last_round:
         rewind_weights_functions(model)
-    else:
+    elif not last_round:
         post_round_functions(model)
 
 

src/pquant/core/torch/layers.py

@@ -1404,11 +1404,12 @@ def apply_final_compression(module):
 
 
 def call_post_round_functions(model, rewind, rounds, r):
+    last_round = (r == rounds - 1)
     if rewind == "round":
         rewind_weights_functions(model)
-    elif rewind == "post-ticket-search" and r == rounds - 1:
+    elif rewind == "post-ticket-search" and last_round:
         rewind_weights_functions(model)
-    elif r != rounds - 1:
+    elif not last_round:
         post_round_functions(model)