ensure resencode to be int64 (#26)

Co-authored-by: Pegerto Fernandez <[email protected]>

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "graphpro"
-version = "0.9.1"
+version = "0.9.2"
 authors = [
   { name="Pegerto Fernandez", email="[email protected]" },
 ]

src/graphpro/annotations.py

@@ -38,4 +38,4 @@ def generate(self, G: Graph, atom_group: AtomGroup):
     def encode(self, G: Graph) -> torch.tensor:
         res_names = [G.node_attr(n)['resname'] for n in G.nodes()]
         res_ids = [self.res_letters.index(name) for name in res_names]
-        return F.one_hot(torch.tensor(res_ids), num_classes=len(self.res_letters))
+        return F.one_hot(torch.tensor(res_ids, dtype=torch.int64), num_classes=len(self.res_letters))

test/graphpro/annotations_test.py

@@ -1,3 +1,4 @@
+import torch
 import MDAnalysis as mda
 
 from graphpro import md_analisys
@@ -19,4 +20,5 @@ def test_resname_encoded():
     G = md_analisys(u1).generate(ContactMap(cutoff=6), [ResidueType()])
     data = G.to_data(node_encoders=  [ResidueType()])
 
-    assert data.x.size() == (214, 22)
+    assert data.x.size() == (214, 22)
+    assert data.x.dtype == torch.int64