Skip to content

GitLab

Unverified Commit 4f30aae0 authored by Rocco Meli's avatar Rocco Meli Committed by GitHub
Browse files

fix #461 (#462)

parent b12e7e6e
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
torchani/nn.py
View file @ 4f30aae0
...@@ -123,4 +123,4 @@ class SpeciesConverter(torch.nn.Module): ...@@ -123,4 +123,4 @@ class SpeciesConverter(torch.nn.Module):
pbc: Optional[Tensor] = None): pbc: Optional[Tensor] = None):
"""Convert species from periodic table element index to 0, 1, 2, 3, ... indexing""" """Convert species from periodic table element index to 0, 1, 2, 3, ... indexing"""
species, coordinates = input_ species, coordinates = input_
return SpeciesCoordinates(self.conv_tensor[species], coordinates) return SpeciesCoordinates(self.conv_tensor[species].to(species.device), coordinates)
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