linting

pyg-team · Mar 21, 2022 · 9026ebc · 9026ebc
1 parent c910e96
commit 9026ebc
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/torch_geometric/nn/conv/pna_conv.py b/torch_geometric/nn/conv/pna_conv.py
@@ -105,14 +105,12 @@ def __init__(self, in_channels: int, out_channels: int,
         self.F_out = self.out_channels // towers
 
         deg = deg.to(torch.float)
-        total_no_vertices = deg.sum()
-        bin_degrees = torch.arange(len(deg))
+        num_nodes = int(deg.sum())
+        bin_degrees = torch.arange(deg.numel())
         self.avg_deg: Dict[str, float] = {
-            'lin': ((bin_degrees * deg).sum() / total_no_vertices).item(),
-            'log':
-            (((bin_degrees + 1).log() * deg).sum() / total_no_vertices).item(),
-            'exp':
-            ((bin_degrees.exp() * deg).sum() / total_no_vertices).item(),
+            'lin': float((bin_degrees * deg).sum()) / num_nodes,
+            'log': float(((bin_degrees + 1).log() * deg).sum()) / num_nodes,
+            'exp': float((bin_degrees.exp() * deg).sum()) / num_nodes,
         }
 
         if self.edge_dim is not None: