fix(normal_forward_compute_actor): fix SendMsgToForwardModelSaveActor() (#1270)

* fix(normal_forward_compute_actor): fix SendMsgToForwardModelSaveActor()

* refine(normal_forward_compute_actor)

oneflow/core/actor/actor.h

@@ -89,6 +89,7 @@ class Actor {
 
   // Util For Derived Actor to Send Msg
   void AsyncSendMsg(const ActorMsg&);
+  int64_t HandleRegstToConsumer(Regst* regst, std::function<bool(int64_t)> IsAllowedActor);
   void HandleProducedNaiveDataRegstToConsumer(std::function<bool(Regst*)> RegstPreProcess,
                                               std::function<bool(int64_t)> IsAllowedActor);
   void HandleProducedNaiveDataRegstToConsumer(std::function<bool(Regst*)> RegstPreProcess);
@@ -170,7 +171,6 @@ class Actor {
   virtual void AsyncSendCustomizedConsumedRegstMsgToProducer() {}
   void AsyncSendConsumedCtrlRegstMsgToProducer();
   void AsyncSendProducedCtrlRegstMsgToConsumer();
-  int64_t HandleRegstToConsumer(Regst* regst, std::function<bool(int64_t)> IsAllowedActor);
   virtual void AsyncReturnAllCustomizedReadableRegst() {}
 
   int64_t actor_id_;

oneflow/core/actor/normal_forward_compute_actor.cpp

@@ -203,10 +203,10 @@ void NormalForwardCompActor::TrySendMsgToForwardModelSaveActor(int64_t piece_id)
 }
 
 void NormalForwardCompActor::SendMsgToForwardModelSaveActor(int64_t batch_id) {
-  HandleProducedNaiveDataRegstToConsumer([&](Regst* regst) {
-    regst->set_model_version_id(batch_id);
-    return regst->regst_desc_id() == forward_model_regst_desc_id_;
-  });
+  Regst* fw_model_regst = GetNaiveCurWriteable(forward_model_regst_desc_id_);
+  CHECK(fw_model_regst);
+  fw_model_regst->set_model_version_id(batch_id);
+  HandleRegstToConsumer(fw_model_regst, [](int64_t) { return true; });
 }
 
 void NormalForwardCompActor::SendConstBufInitMsgToBwActor() {