diff --git a/containers/localnet/change_genesis.py b/containers/localnet/change_genesis.py
index fc4b544..8115c02 100644
--- a/containers/localnet/change_genesis.py
+++ b/containers/localnet/change_genesis.py
@@ -172,7 +172,6 @@ def replace_in_files():
         with open(external_json_path, 'r') as f:
             existing_data = json.load(f)
 
-        print(f"validators_objects: {validators_objects}, shards_count: {shards_count}")
         existing_data = validators_objects[0:shards_count] + existing_data
 
         existing_data.append(genesis_object)
diff --git a/containers/localnet/run.sh b/containers/localnet/run.sh
index ba29029..8299774 100644
--- a/containers/localnet/run.sh
+++ b/containers/localnet/run.sh
@@ -55,7 +55,7 @@ setup() {
   cd localnet && sudo python3 replace_economics.py "$MX_RESULT_TOTAL_SUPPLY" && cd ..
 
   echo "Applying workaround for enableEpochs..."
-  cd localnet && sudo python3 temp_replace_enable_epochs.py && cd ..
+  cd localnet && sudo python3 temp_replace_enable_epochs.py "$MX_LT_NUM_SHARDS" && cd ..
 
   echo "Applying workaround for genesisContracts..."
   cd localnet && sudo python3 temp_replace_genesis_smart_contracts.py && cd ..
diff --git a/containers/localnet/temp_replace_enable_epochs.py b/containers/localnet/temp_replace_enable_epochs.py
index 2b1fda5..d194982 100644
--- a/containers/localnet/temp_replace_enable_epochs.py
+++ b/containers/localnet/temp_replace_enable_epochs.py
@@ -1,9 +1,13 @@
 import os
 import re
 import subprocess
+import sys
 
+num_shards = sys.argv[1]
+staking_v4_step3_max_num_nodes = 64 - (int(num_shards) + 1) * 2
 
-enable_epochs_content = """
+
+enable_epochs_content = f"""
     [EnableEpochs]
     # SCDeployEnableEpoch represents the epoch when the deployment of smart contracts will be enabled
     SCDeployEnableEpoch = 1
@@ -295,8 +299,8 @@
 
     # BLSMultiSignerEnableEpoch represents the activation epoch for different types of BLS multi-signers
     BLSMultiSignerEnableEpoch = [
-        { EnableEpoch = 0, Type = "no-KOSK" },
-        { EnableEpoch = 1, Type = "KOSK" }
+        {{ EnableEpoch = 0, Type = "no-KOSK" }},
+        {{ EnableEpoch = 1, Type = "KOSK" }}
     ]
 
     # StakeLimitsEnableEpoch represents the epoch when stake limits on validators are enabled
@@ -315,19 +319,19 @@
 
     # MaxNodesChangeEnableEpoch holds configuration for changing the maximum number of nodes and the enabling epoch
     MaxNodesChangeEnableEpoch = [
-        { EpochEnable = 0, MaxNumNodes = 48, NodesToShufflePerShard = 4 },  # 4 shuffled out keys / shard will not be reached normally
-        { EpochEnable = 1, MaxNumNodes = 64, NodesToShufflePerShard = 2 },
+        {{ EpochEnable = 0, MaxNumNodes = 48, NodesToShufflePerShard = 4 }},  # 4 shuffled out keys / shard will not be reached normally
+        {{ EpochEnable = 1, MaxNumNodes = 64, NodesToShufflePerShard = 2 }},
         # Staking v4 configuration, where:
         # - Enable epoch = StakingV4Step3EnableEpoch
         # - NodesToShufflePerShard = same as previous entry in MaxNodesChangeEnableEpoch
         # - MaxNumNodes = (MaxNumNodesFromPreviousEpochEnable - (numOfShards+1)*NodesToShufflePerShard)
-        { EpochEnable = 6, MaxNumNodes = 60, NodesToShufflePerShard = 2 },
+        {{ EpochEnable = 6, MaxNumNodes = {staking_v4_step3_max_num_nodes}, NodesToShufflePerShard = 2 }},
     ]
 
     [GasSchedule]
     # GasScheduleByEpochs holds the configuration for the gas schedule that will be applied from specific epochs
     GasScheduleByEpochs = [
-        { StartEpoch = 0, FileName = "gasScheduleV7.toml" },
+        {{ StartEpoch = 0, FileName = "gasScheduleV7.toml" }},
     ]
 """