fix #461 (#462)

4f30aae0 · Rocco Meli · GitHub · b12e7e6e · 4f30aae0
Unverified Commit 4f30aae0 authored Apr 29, 2020 by Rocco Meli Committed by GitHub Apr 29, 2020
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

torchani/nn.py torchani/nn.py +1 -1

No files found.
--- a/torchani/nn.py
+++ b/torchani/nn.py
@@ -123,4 +123,4 @@ class SpeciesConverter(torch.nn.Module):
                pbc: Optional[Tensor] = None):
        """Convert species from periodic table element index to 0, 1, 2, 3, ... indexing"""
        species, coordinates = input_
-        return SpeciesCoordinates(self.conv_tensor[species], coordinates)
+        return SpeciesCoordinates(self.conv_tensor[species].to(species.device), coordinates)