Use searchsorted over argmax (#9670)

Lightning-AI · Sep 23, 2021 · 2b2537d · 2b2537d
1 parent 568a1e0
commit 2b2537d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/loops/batch/training_batch_loop.py b/pytorch_lightning/loops/batch/training_batch_loop.py
@@ -191,5 +191,5 @@ def get_active_optimizers(self, batch_idx: Optional[int] = None) -> List[Tuple[i
         current_place_in_loop = batch_idx % optimizers_loop_length
 
         # find optimzier index by looking for the first {item > current_place} in the cumsum list
-        opt_idx = int(np.argmax(self.optimizer_freq_cumsum > current_place_in_loop))
+        opt_idx = np.searchsorted(self.optimizer_freq_cumsum, current_place_in_loop, side="right")
         return [(opt_idx, self.trainer.optimizers[opt_idx])]