diff --git a/python/bigdl/dllib/optim/optimizer.py b/python/bigdl/dllib/optim/optimizer.py
index <HASH>..<HASH> 100644
--- a/python/bigdl/dllib/optim/optimizer.py
+++ b/python/bigdl/dllib/optim/optimizer.py
@@ -630,6 +630,13 @@ class BaseOptimizer(JavaValue):
         print("Loading input ...")
         self.value.prepareInput()
 
+    def set_end_when(self, end_when):
+        """
+        When to stop, passed in a [[Trigger]]
+        """
+        self.value.setEndWhen(end_when.value)
+        return self
+
 
 class Optimizer(BaseOptimizer):
Fix optimizer state messed up when calling optimize() multiple times in DistriOptimizer (#<I>)

* support continue training



* meet code review



* add unit tests



* fix set model