Fix typing for DGL 2.0 (#79)

* use dglteam dgl

* set node type

* fix astype

* type edges

* update changelog

CHANGELOG.md

@@ -18,13 +18,15 @@ The rules for this file:
 
 ### Authors
 <!-- GitHub usernames of contributors to this release -->
+- lilyminium
 
 ### Reviewers
 
 ### Added
 
 ### Fixed
 <!-- Bug fixes -->
+- Fixed node and edge typing, adding DGL 2.0 compatibility (Issue #78, PR #79)
 
 ### Changed
 <!-- Changes in existing functionality -->

devtools/conda-envs/test_env.yaml

@@ -1,6 +1,7 @@
 name: openff-nagl-test
 channels:
   - openeye
+  - dglteam
   - conda-forge
   - defaults
 dependencies:

openff/nagl/molecule/_dgl/molecule.py

@@ -97,13 +97,13 @@ def from_openff(
             for offmol in offmols
         ]
         graph = dgl.batch(subgraphs)
-        graph.set_batch_num_nodes(graph.batch_num_nodes().sum().reshape((-1,)))
-        graph.set_batch_num_edges(
-            {
-                e_type: graph.batch_num_edges(e_type).sum().reshape((-1,))
-                for e_type in graph.canonical_etypes
-            }
-        )
+        n_nodes = graph.batch_num_nodes().sum().reshape((-1,))
+        graph.set_batch_num_nodes(n_nodes.type(torch.int32))
+        edges = {}
+        for e_type in graph.canonical_etypes:
+            n_edge = graph.batch_num_edges(e_type).sum().reshape((-1,))
+            edges[e_type] = n_edge.type(torch.int32)
+        graph.set_batch_num_edges(edges)
 
         mapped_smiles = offmols[0].to_smiles(mapped=True)