[Misc] Black auto fix. (dmlc#4651)

Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>

Author: frozenbugs

examples/pytorch/GNN-FiLM/data_loader.py

@@ -1,61 +1,98 @@
-from torch.utils.data import Dataset, DataLoader
+import collections
+
+from torch.utils.data import DataLoader, Dataset
+
 import dgl
 from dgl.data import PPIDataset
-import collections
 
-#implement the collate_fn for dgl graph data class
-PPIBatch = collections.namedtuple('PPIBatch', ['graph', 'label'])
+# implement the collate_fn for dgl graph data class
+PPIBatch = collections.namedtuple("PPIBatch", ["graph", "label"])
+
+
 def batcher(device):
     def batcher_dev(batch):
         batch_graphs = dgl.batch(batch)
-        return PPIBatch(graph=batch_graphs,
-                        label=batch_graphs.ndata['label'].to(device))
+        return PPIBatch(
+            graph=batch_graphs, label=batch_graphs.ndata["label"].to(device)
+        )
+
     return batcher_dev
 
-#add a fresh "self-loop" edge type to the untyped PPI dataset and prepare train, val, test loaders
-def load_PPI(batch_size=1, device='cpu'):
-    train_set = PPIDataset(mode='train')
-    valid_set = PPIDataset(mode='valid')
-    test_set = PPIDataset(mode='test')
-    #for each graph, add self-loops as a new relation type
-    #here we reconstruct the graph since the schema of a heterograph cannot be changed once constructed
+
+# add a fresh "self-loop" edge type to the untyped PPI dataset and prepare train, val, test loaders
+def load_PPI(batch_size=1, device="cpu"):
+    train_set = PPIDataset(mode="train")
+    valid_set = PPIDataset(mode="valid")
+    test_set = PPIDataset(mode="test")
+    # for each graph, add self-loops as a new relation type
+    # here we reconstruct the graph since the schema of a heterograph cannot be changed once constructed
     for i in range(len(train_set)):
-        g = dgl.heterograph({
-            ('_N','_E','_N'): train_set[i].edges(),
-            ('_N', 'self', '_N'): (train_set[i].nodes(), train_set[i].nodes())
-        })
-        g.ndata['label'] = train_set[i].ndata['label']
-        g.ndata['feat'] = train_set[i].ndata['feat']
-        g.ndata['_ID'] = train_set[i].ndata['_ID']
-        g.edges['_E'].data['_ID'] = train_set[i].edata['_ID']
+        g = dgl.heterograph(
+            {
+                ("_N", "_E", "_N"): train_set[i].edges(),
+                ("_N", "self", "_N"): (
+                    train_set[i].nodes(),
+                    train_set[i].nodes(),
+                ),
+            }
+        )
+        g.ndata["label"] = train_set[i].ndata["label"]
+        g.ndata["feat"] = train_set[i].ndata["feat"]
+        g.ndata["_ID"] = train_set[i].ndata["_ID"]
+        g.edges["_E"].data["_ID"] = train_set[i].edata["_ID"]
         train_set.graphs[i] = g
     for i in range(len(valid_set)):
-        g = dgl.heterograph({
-            ('_N','_E','_N'): valid_set[i].edges(),
-            ('_N', 'self', '_N'): (valid_set[i].nodes(), valid_set[i].nodes())
-        })
-        g.ndata['label'] = valid_set[i].ndata['label']
-        g.ndata['feat'] = valid_set[i].ndata['feat']
-        g.ndata['_ID'] = valid_set[i].ndata['_ID']
-        g.edges['_E'].data['_ID'] = valid_set[i].edata['_ID']
-        valid_set.graphs[i] = g    
+        g = dgl.heterograph(
+            {
+                ("_N", "_E", "_N"): valid_set[i].edges(),
+                ("_N", "self", "_N"): (
+                    valid_set[i].nodes(),
+                    valid_set[i].nodes(),
+                ),
+            }
+        )
+        g.ndata["label"] = valid_set[i].ndata["label"]
+        g.ndata["feat"] = valid_set[i].ndata["feat"]
+        g.ndata["_ID"] = valid_set[i].ndata["_ID"]
+        g.edges["_E"].data["_ID"] = valid_set[i].edata["_ID"]
+        valid_set.graphs[i] = g
     for i in range(len(test_set)):
-        g = dgl.heterograph({
-            ('_N','_E','_N'): test_set[i].edges(),
-            ('_N', 'self', '_N'): (test_set[i].nodes(), test_set[i].nodes())
-        })
-        g.ndata['label'] = test_set[i].ndata['label']
-        g.ndata['feat'] = test_set[i].ndata['feat']
-        g.ndata['_ID'] = test_set[i].ndata['_ID']
-        g.edges['_E'].data['_ID'] = test_set[i].edata['_ID']
-        test_set.graphs[i] = g  
+        g = dgl.heterograph(
+            {
+                ("_N", "_E", "_N"): test_set[i].edges(),
+                ("_N", "self", "_N"): (
+                    test_set[i].nodes(),
+                    test_set[i].nodes(),
+                ),
+            }
+        )
+        g.ndata["label"] = test_set[i].ndata["label"]
+        g.ndata["feat"] = test_set[i].ndata["feat"]
+        g.ndata["_ID"] = test_set[i].ndata["_ID"]
+        g.edges["_E"].data["_ID"] = test_set[i].edata["_ID"]
+        test_set.graphs[i] = g
 
     etypes = train_set[0].etypes
-    in_size = train_set[0].ndata['feat'].shape[1]
-    out_size = train_set[0].ndata['label'].shape[1]
+    in_size = train_set[0].ndata["feat"].shape[1]
+    out_size = train_set[0].ndata["label"].shape[1]
 
-    #prepare train, valid, and test dataloaders
-    train_loader = DataLoader(train_set, batch_size=batch_size, collate_fn=batcher(device), shuffle=True)
-    valid_loader = DataLoader(valid_set, batch_size=batch_size, collate_fn=batcher(device), shuffle=True)
-    test_loader = DataLoader(test_set, batch_size=batch_size, collate_fn=batcher(device), shuffle=True)
+    # prepare train, valid, and test dataloaders
+    train_loader = DataLoader(
+        train_set,
+        batch_size=batch_size,
+        collate_fn=batcher(device),
+        shuffle=True,
+    )
+    valid_loader = DataLoader(
+        valid_set,
+        batch_size=batch_size,
+        collate_fn=batcher(device),
+        shuffle=True,
+    )
+    test_loader = DataLoader(
+        test_set,
+        batch_size=batch_size,
+        collate_fn=batcher(device),
+        shuffle=True,
+    )
     return train_loader, valid_loader, test_loader, etypes, in_size, out_size