Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Feb 3, 2024
1 parent 6750312 commit 0a3e208
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 9 deletions.
2 changes: 1 addition & 1 deletion deepmd/pt/model/model/dp_atomic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def get_rcut(self) -> float:
def get_sel(self) -> List[int]:
"""Get the neighbor selection."""
return self.sel

def get_nnei(self) -> List[int]:
"""Get the the number of selected atoms in the cut-off radius."""
return self.nnei

Check warning on line 106 in deepmd/pt/model/model/dp_atomic_model.py

View check run for this annotation

Codecov / codecov/patch

deepmd/pt/model/model/dp_atomic_model.py#L106

Added line #L106 was not covered by tests
Expand Down
22 changes: 14 additions & 8 deletions source/api_cc/src/DeepPotPT.cc
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,11 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
const std::vector<VALUETYPE>& box,
const InputNlist& lmp_list,
const int& ago) {
if(cpu_enabled)
if (cpu_enabled) {
torch::Device device(torch::kCPU);
else
} else {
torch::Device device(torch::kCUDA, gpu_id);
}
std::vector<VALUETYPE> coord_wrapped = coord;
int natoms = atype.size();
auto options = torch::TensorOptions().dtype(torch::kFloat64);
Expand All @@ -93,9 +94,10 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
at::Tensor firstneigh = torch::from_blob(
nlist_data.jlist, {lmp_list.inum, max_num_neighbors}, int32_options);
at::Tensor nlist = firstneigh.to(torch::kInt64).to(device);
firstneigh_tensor =
module.run_method("format_nlist", coord_wrapped_Tensor,atype_Tensor,nlist)
.toTensor();
firstneigh_tensor = module
.run_method("format_nlist", coord_wrapped_Tensor,
atype_Tensor, nlist)
.toTensor();
} else {
at::Tensor firstneigh = torch::from_blob(
nlist_data.jlist, {1, lmp_list.inum, max_num_neighbors},
Expand All @@ -106,7 +108,10 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
bool do_atom_virial_tensor = true;
torch::Tensor mapping_tensor = torch::Tensor(nullptr);
c10::Dict<c10::IValue, c10::IValue> outputs =
module.run_method("forward_lower",coord_wrapped_Tensor,atype_Tensor,firstneigh_tensor,mapping_tensor, do_atom_virial_tensor).toGenericDict();
module
.run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor,
firstneigh_tensor, mapping_tensor, do_atom_virial_tensor)
.toGenericDict();
c10::IValue energy_ = outputs.at("energy");
c10::IValue force_ = outputs.at("extended_force");
c10::IValue virial_ = outputs.at("reduced_virial");
Expand Down Expand Up @@ -190,10 +195,11 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
const std::vector<VALUETYPE>& coord,
const std::vector<int>& atype,
const std::vector<VALUETYPE>& box) {
if(cpu_enabled)
if (cpu_enabled) {
torch::Device device(torch::kCPU);
else
} else {
torch::Device device(torch::kCUDA, gpu_id);
}
std::vector<VALUETYPE> coord_wrapped = coord;
int natoms = atype.size();
auto options = torch::TensorOptions().dtype(torch::kFloat64);
Expand Down

0 comments on commit 0a3e208

Please sign in to comment.