Update save_pruned_model.py

lucas-ecm · Feb 11, 2023 · af65c67 · af65c67
1 parent c5a4a4d
commit af65c67
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/save_pruned_model.py b/save_pruned_model.py
@@ -60,8 +60,9 @@ def save_tf():
   epochs = 2
   validation_split = 0.1 # 10% of training set will be used for validation set. 
 
-  num_images = train_images.shape[0] * (1 - validation_split)
-  end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs
+  #num_images = train_images.shape[0] * (1 - validation_split)
+  #end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs
+  end_step = 1000
 
   # Define model for pruning.
   pruning_params = {