Method Forward

src/model/layer/flatten.cc:37–47

}

const Tensor Flatten::Forward(int flag, const Tensor &input) {
  Tensor output;
  input_shape_ = input.shape();
  if (axis_ == 0)
    output = Reshape(input, vector<size_t>{input.Size()});
  else
    output =
        Reshape(input, vector<size_t>{input.Size() / out_sample_shape_.at(0),
                                      out_sample_shape_.at(0)});
  return output;
}

const std::pair<Tensor, vector<Tensor> > Flatten::Backward(int flag,
                                                           const Tensor &grad) {