Skip to content

GitLab

Unverified Commit 566320ec authored by Farhad Ramezanghorbani's avatar Farhad Ramezanghorbani Committed by GitHub
Browse files

Fix CUDA support in AEVComputer (#350) 

* fix cuda support

* Update aev.py
parent 38573358
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torchani/aev.py
View file @ 566320ec
...@@ -362,7 +362,7 @@ class AEVComputer(torch.nn.Module): ...@@ -362,7 +362,7 @@ class AEVComputer(torch.nn.Module):
self.aev_length = self.radial_length + self.angular_length self.aev_length = self.radial_length + self.angular_length
self.sizes = self.num_species, self.radial_sublength, self.radial_length, self.angular_sublength, self.angular_length, self.aev_length self.sizes = self.num_species, self.radial_sublength, self.radial_length, self.angular_sublength, self.angular_length, self.aev_length
self.register_buffer('triu_index', triu_index(num_species)) self.register_buffer('triu_index', triu_index(num_species).to(device=self.EtaR.device))
# Set up default cell and compute default shifts. # Set up default cell and compute default shifts.
# These values are used when cell and pbc switch are not given. # These values are used when cell and pbc switch are not given.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment