Merge pull request #4 from VibhuJawa/cugraph-dgl-dlfw-patch

Fix dlpack boolean errors

python/cugraph-dgl/examples/graphsage/node-classification.py

@@ -9,6 +9,9 @@
 
 # Timing Imports
 import time
+# Ignore Warning
+import warnings
+warnings.filterwarnings('ignore')
 
 import torch
 import torch.nn as nn
@@ -216,6 +219,16 @@ def train(args, device, g, dataset, model):
                 pool_allocator=True, initial_pool_size=5e9, maximum_pool_size=25e9
             )
 
+        # Work around for DLFW container issues
+        # where dlpack conversion of boolean fails 
+        # on the first run
+        # Similar to issue https://github.com/dmlc/dgl/issues/3591
+        if 'train_mask' in g.ndata:
+            g.ndata['train_mask'] = g.ndata['train_mask'].int()
+        if 'train_mask' in g.ndata:
+            g.ndata['test_mask'] = g.ndata['test_mask'].int()
+        if 'val_mask' in g.ndata: 
+            g.ndata['val_mask'] = g.ndata['val_mask'].int()
         g = cugraph_dgl.cugraph_storage_from_heterograph(g)
         del dataset.g