update to follow schnet major update

OMalenfantThuot · Apr 30, 2021 · bcbbe5f · bcbbe5f
1 parent 5d0099c
commit bcbbe5f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/mlcalcdriver/interfaces/schnetpack_interface.py b/mlcalcdriver/interfaces/schnetpack_interface.py
@@ -1,7 +1,7 @@
 import torch
 import numpy as np
 from torch.utils.data import Dataset
-from schnetpack.data.atoms import _convert_atoms
+from schnetpack.data.atoms import _convert_atoms, torchify_dict
 
 
 class SchnetPackData(Dataset):
@@ -26,8 +26,8 @@ def __getitem__(self, idx):
         Needed to create a PyTorch Dataset
         """
         _, properties = self.get_properties(idx)
-        properties["_idx"] = torch.LongTensor(np.array([idx], dtype=np.int))
-        return properties
+        properties["_idx"] = np.array([idx], dtype=np.int)
+        return torchify_dict(properties)
 
     def get_properties(self, idx):
         r"""

diff --git a/requirements_dev.txt b/requirements_dev.txt
@@ -11,7 +11,7 @@ pytest-cov==2.11.1
 pytest-sugar==0.9.4
 coveralls==3.0.1
 ase==3.21.1
-git+https://github.com/atomistic-machine-learning/schnetpack.git@b5e04d43322809ba308d3c1a66e916aa617e3370
+git+https://github.com/atomistic-machine-learning/schnetpack.git
 sphinx==3.5.4
 sphinx-rtd-theme==0.5.2
 git+https://github.com/crossnox/m2r@dev#egg=m2r