Revert commit c96641e - num_replicas parameter is enough

Author: ASvyatkovskiy

plasma/models/mpi_runner.py

@@ -139,7 +139,7 @@ def get_val(self):
 
 
 class MPIModel():
-  def __init__(self,model,optimizer,comm,batch_iterator,batch_size,num_replicas=None,warmup_steps=1000,lr=0.01,custom_num_workers=0):
+  def __init__(self,model,optimizer,comm,batch_iterator,batch_size,num_replicas=None,warmup_steps=1000,lr=0.01):
     # random.seed(task_index)
     self.epoch = 0
     self.model = model
@@ -151,19 +151,13 @@ def __init__(self,model,optimizer,comm,batch_iterator,batch_size,num_replicas=No
     self.batch_size = batch_size
     self.batch_iterator = batch_iterator
     self.warmup_steps=warmup_steps
-    if custom_num_workers:
-        if custom_num_workers < comm.Get_size():
-            self.num_workers = custom_num_workers
-        else: self.num_workers = comm.Get_size()
-    else:
-        self.num_workers = comm.Get_size()
+    self.num_workers = comm.Get_size()
     self.task_index = comm.Get_rank()
     self.history = cbks.History()
     if num_replicas is None or num_replicas < 1 or num_replicas > self.num_workers:
-      self.num_replicas = self.num_workers
+        self.num_replicas = self.num_workers
     else:
-      self.num_replicas = num_replicas
-
+        self.num_replicas = num_replicas
 
   def set_lr(self,lr):
     self.lr = lr