Browse Source

fix weight_update_mod bug

tags/v1.2.0
xutianchun 5 years ago
parent
commit
d365f0a296
1 changed files with 3 additions and 3 deletions
  1. +3
    -3
      mindspore/lite/src/train/optimizer_kernel.h

+ 3
- 3
mindspore/lite/src/train/optimizer_kernel.h View File

@@ -32,7 +32,7 @@ class OptimizerKernel : public LiteKernel {
~OptimizerKernel() = default;

enum class WeightUpdateMode { NORMAL, VIRTUAL_BATCH };
WeightUpdateMode get_optimizer_mode() { return weightUpdateMod_; }
WeightUpdateMode get_optimizer_mode() { return weight_update_mod_; }

int Init() override {
default_lr_ = reinterpret_cast<float *>(in_tensors_.at(lr_idx_)->MutableData())[0];
@@ -67,6 +67,7 @@ class OptimizerKernel : public LiteKernel {
}
valid_grad_sum_ = false;
std::fill(grad_sum_, grad_sum_ + elem_num, 0);
weight_update_mod_ = WeightUpdateMode::VIRTUAL_BATCH;
} else {
if (grad_sum_ != nullptr) {
OptimizerStep();
@@ -74,7 +75,6 @@ class OptimizerKernel : public LiteKernel {
grad_sum_ = nullptr;
}
}
weightUpdateMod_ = WeightUpdateMode::VIRTUAL_BATCH;
return RET_OK;
}

@@ -112,7 +112,7 @@ class OptimizerKernel : public LiteKernel {
bool valid_grad_sum_ = false;

private:
WeightUpdateMode weightUpdateMod_ = WeightUpdateMode::NORMAL;
WeightUpdateMode weight_update_mod_ = WeightUpdateMode::NORMAL;
};

} // namespace mindspore::kernel


Loading…
Cancel
Save