[Bug] Hardgat enable gpu training (#2632)

* add multihead in DotGatConv

* Fix spacing issue

* Add Unit test for dotgat

* Modified Unit test for dotgat

* Add transformer like divisor

* Update dotgatconv.py

* Update hgao.py

* Update train.py
Co-authored-by: Chen <chesirui@3c22fbe5458c.ant.amazon.com>
Co-authored-by: Zihao Ye <expye@outlook.com>

examples/pytorch/hardgat/hgao.py

@@ -79,7 +79,7 @@ class HardGAO(nn.Module):
             # Use edge message passing function to get the weight from src node
             graph.apply_edges(fn.copy_u('y','y'))
             # Select Top k neighbors
-            subgraph = select_topk(graph,self.k,'y')
+            subgraph = select_topk(graph.cpu(),self.k,'y').to(graph.device)
             # Sigmoid as information threshold
             subgraph.ndata['y'] = torch.sigmoid(subgraph.ndata['y'])
             # Using vector matrix elementwise mul for acceleration

examples/pytorch/hardgat/train.py

@@ -51,7 +51,7 @@ def main(args):
         cuda = False
     else:
         cuda = True
-        g = g.int().to(args.gpu)
+        g = g.to(args.gpu)
 
     features = g.ndata['feat']
     labels = g.ndata['label']