fix dump field (PaddlePaddle#67)

wangzhen38 · Oct 26, 2022 · cff4bea · cff4bea
1 parent 0e31882
commit cff4bea
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/paddle/fluid/framework/data_feed.h b/paddle/fluid/framework/data_feed.h
@@ -928,6 +928,7 @@ class DataFeed {
     return ins_content_vec_;
   }
   virtual int GetCurBatchSize() { return batch_size_; }
+  virtual void SetCurBatchSize(const int batch_size) { batch_size_ = batch_size; }
   virtual void LoadIntoMemory() {
     PADDLE_THROW(platform::errors::Unimplemented(
         "This function(LoadIntoMemory) is not implemented."));

diff --git a/paddle/fluid/framework/ps_gpu_worker.cc b/paddle/fluid/framework/ps_gpu_worker.cc
@@ -458,7 +458,7 @@ void PSGPUWorker::TrainFiles() {
     if (cur_batch <= 0) {
       break;
     }
-
+    device_reader_->SetCurBatchSize(cur_batch);
     total_ins_num += cur_batch;
 
     if (shape_check_flag_.load()) {