diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizerV2.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizerV2.scala
index ec20e4646d4..e8c9b5edb33 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizerV2.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/DistriOptimizerV2.scala
@@ -56,16 +56,7 @@ object DistriOptimizerV2 extends AbstractOptimizer {
     parameterProcessers: Array[ParameterProcessor] = null) extends DistriOptimizer.Cache[T]
 
   import Optimizer._
-
-  private[DistriOptimizerV2] var _logger: Option[OptimizerLogger] = None
-
-  def logger: OptimizerLogger = {
-    if (_logger.isEmpty) {
-      _logger = Some(new DistriLogger)
-    }
-
-    _logger.get
-  }
+  val logger: Logger = Logger.getLogger(getClass)
 
   private[optim] def optimize[T: ClassTag](
     cacheOfMaster: MasterCache[T],
@@ -107,7 +98,7 @@ object DistriOptimizerV2 extends AbstractOptimizer {
       }
 
       val _header = header(
-        trainingTrace.epochs,
+        trainingTrace.epochs-1,
         trainingTrace.recordsOfEpoch,
         context.numSamples,
         trainingTrace.iterations,
@@ -857,10 +848,6 @@ class DistriOptimizerV2[T: ClassTag](
     CachedModels.deleteKey(modelBroadcast.uuid)
   }
 
-  def setLogger(logger: OptimizerLogger): Unit = {
-    DistriOptimizerV2._logger = Some(logger)
-  }
-
   private def validArgs(): Boolean = {
     val checkSingleton = this.checkSingleton
     val nodeNumber = Engine.nodeNumber()
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LoggerFilterSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LoggerFilterSpec.scala
index ed2363ae57a..1a144882873 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LoggerFilterSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LoggerFilterSpec.scala
@@ -121,7 +121,7 @@ class LoggerFilterSpec extends FlatSpec with BeforeAndAfter with Matchers {
     }
 
     {
-      val pattern = s".*INFO.*DistriLogger.* - " + "" +
+      val pattern = s".*INFO.*DistriOptimizerV2.* - " + "" +
         s"\\[Epoch 1 100/100\\]\\[Iteration 2\\]\\[Wall Clock .*\\] " +
         s"Epoch finished. Wall clock time is .*ms"