diff --git a/bajor/env_helpers.py b/bajor/env_helpers.py
index 7faaaa6..60ccf9d 100644
--- a/bajor/env_helpers.py
+++ b/bajor/env_helpers.py
@@ -16,9 +16,9 @@ def port():
     return os.environ.get('PORT', '8000')
 
 def max_num_pool_nodes(pool_id):
-    if(pool_id == 'training_1'):
+    if('training' in pool_id):
         return os.environ.get('MAX_NODES_TRAINING', 2)
-    elif(pool_id == 'predictions_0'):
+    elif('predictions' in pool_id):
         return os.environ.get('MAX_NODES_PREDICTION', 2)
     else:
         return 0