[Cleanup] Change Byte to Bool for training masks (#954)

* Change Byte to Bool for training masks

* Check if module has Bool, otherwise use Byte

examples/pytorch/appnp/train.py

@@ -25,6 +25,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/cluster_gcn/cluster_gcn.py

@@ -46,9 +46,14 @@ def main(args):
         labels = torch.LongTensor(data.labels)
     else:
         labels = torch.FloatTensor(data.labels)
-    train_mask = torch.ByteTensor(data.train_mask).type(torch.bool)
-    val_mask = torch.ByteTensor(data.val_mask).type(torch.bool)
-    test_mask = torch.ByteTensor(data.test_mask).type(torch.bool)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
+        train_mask = torch.ByteTensor(data.train_mask)
+        val_mask = torch.ByteTensor(data.val_mask)
+        test_mask = torch.ByteTensor(data.test_mask)
     in_feats = features.shape[1]
     n_classes = data.num_labels
     n_edges = data.graph.number_of_edges()

examples/pytorch/dgi/train.py

@@ -22,6 +22,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/gat/train.py

@@ -41,6 +41,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/gcn/gcn_mp.py

@@ -118,6 +118,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/gcn/train.py

@@ -25,6 +25,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/graphsage/graphsage.py

@@ -60,6 +60,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/model_zoo/citation_network/run.py

@@ -45,6 +45,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/sampling/dis_sampling/gcn_cv_sc_train.py

@@ -24,6 +24,11 @@ def main(args):
 
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/sampling/dis_sampling/gcn_ns_sc_train.py

@@ -25,6 +25,11 @@ def main(args):
 
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/sampling/gcn_cv_sc.py

@@ -149,6 +149,11 @@ def main(args):
 
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/sampling/gcn_ns_sc.py

@@ -120,6 +120,11 @@ def main(args):
 
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/sgc/sgc.py

@@ -30,6 +30,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/sgc/sgc_reddit.py

@@ -33,6 +33,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)

examples/pytorch/tagcn/train.py

@@ -23,6 +23,11 @@ def main(args):
     data = load_data(args)
     features = torch.FloatTensor(data.features)
     labels = torch.LongTensor(data.labels)
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(data.train_mask)
+        val_mask = torch.BoolTensor(data.val_mask)
+        test_mask = torch.BoolTensor(data.test_mask)
+    else:
         train_mask = torch.ByteTensor(data.train_mask)
         val_mask = torch.ByteTensor(data.val_mask)
         test_mask = torch.ByteTensor(data.test_mask)