Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions source/api_cc/include/DeepPotPD.h
Original file line number Diff line number Diff line change
Expand Up @@ -282,12 +282,11 @@ class DeepPotPD : public DeepPotBackend {
* @brief Compute the number of elements in a tensor.
* @param[in] x Tensor x.
**/
int numel(const paddle_infer::Tensor& x) const {
// TODO: There might be a overflow problem here for multiply int numbers.
int ret = 1;
size_t numel(const paddle_infer::Tensor& x) const {
size_t ret = 1;
std::vector<int> x_shape = x.shape();
for (std::size_t i = 0, n = x_shape.size(); i < n; ++i) {
ret *= x_shape[i];
ret *= static_cast<size_t>(x_shape[i]);
}
return ret;
};
Expand Down
20 changes: 10 additions & 10 deletions source/api_cc/src/DeepPotPD.cc
Original file line number Diff line number Diff line change
Expand Up @@ -474,9 +474,9 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
auto energy_ = predictor_fl->GetOutputHandle(output_names.at(1));
auto force_ = predictor_fl->GetOutputHandle(output_names.at(2));
auto virial_ = predictor_fl->GetOutputHandle(output_names.at(4));
int output_energy_size = numel(*energy_);
int output_force_size = numel(*force_);
int output_virial_size = numel(*virial_);
size_t output_energy_size = numel(*energy_);
size_t output_force_size = numel(*force_);
size_t output_virial_size = numel(*virial_);
// output energy
ener.resize(output_energy_size);
energy_->CopyToCpu(ener.data());
Expand Down Expand Up @@ -597,26 +597,26 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
auto force_ = predictor->GetOutputHandle(output_names.at(3));
auto virial_ = predictor->GetOutputHandle(output_names.at(5));

int enery_numel = numel(*energy_);
size_t enery_numel = numel(*energy_);
assert(enery_numel > 0);
ener.resize(enery_numel);
energy_->CopyToCpu(ener.data());

int force_numel = numel(*force_);
size_t force_numel = numel(*force_);
assert(force_numel > 0);
force.resize(force_numel);
force_->CopyToCpu(force.data());

int virial_numel = numel(*virial_);
size_t virial_numel = numel(*virial_);
assert(virial_numel > 0);
virial.resize(virial_numel);
virial_->CopyToCpu(virial.data());

if (atomic) {
auto atom_energy_ = predictor->GetOutputHandle(output_names.at(0));
auto atom_virial_ = predictor->GetOutputHandle(output_names.at(1));
int atom_energy_numel = numel(*atom_energy_);
int atom_virial_numel = numel(*atom_virial_);
size_t atom_energy_numel = numel(*atom_energy_);
size_t atom_virial_numel = numel(*atom_virial_);
assert(atom_energy_numel > 0);
assert(atom_virial_numel > 0);
atom_energy.resize(atom_energy_numel);
Expand Down Expand Up @@ -656,7 +656,7 @@ template void DeepPotPD::compute<float, std::vector<ENERGYTYPE>>(
that need to be postprocessed */
void DeepPotPD::get_type_map(std::string& type_map) {
auto type_map_tensor = predictor->GetOutputHandle("buffer_type_map");
int type_map_size = numel(*type_map_tensor);
size_t type_map_size = numel(*type_map_tensor);

std::vector<int> type_map_arr(type_map_size, 0);
type_map_tensor->CopyToCpu(type_map_arr.data());
Expand All @@ -670,7 +670,7 @@ template <typename BUFFERTYPE>
void DeepPotPD::get_buffer(const std::string& buffer_name,
std::vector<BUFFERTYPE>& buffer_array) {
auto buffer_tensor = predictor->GetOutputHandle(buffer_name);
int buffer_size = numel(*buffer_tensor);
size_t buffer_size = numel(*buffer_tensor);
buffer_array.resize(buffer_size);
buffer_tensor->CopyToCpu(buffer_array.data());
}
Expand Down
Loading