diff --git a/autogl/module/model/dgl/gat_dgl.py b/autogl/module/model/dgl/gat_dgl.py
index e54c63e..e7ecfc3 100644
--- a/autogl/module/model/dgl/gat_dgl.py
+++ b/autogl/module/model/dgl/gat_dgl.py
@@ -78,9 +78,7 @@ class GAT(torch.nn.Module):
         
         for i in range(self.num_layer):
             x = F.dropout(x, p=self.args["dropout"], training=self.training)
-            x = self.convs[i](data, x)
-            # concat
-            x = x.view(-1, self.heads * self.out_channels)
+            x = self.convs[i](data, x).flatten(1)
             if i != self.num_layer - 1:
                 x = activate_func(x, self.args["act"])
 
@@ -89,9 +87,7 @@ class GAT(torch.nn.Module):
     def lp_encode(self, data):
         x = data.ndata['x']
         for i in range(self.num_layer - 1):
-            x = self.convs[i](x, data.train_pos_edge_index)
-            # concat
-            x = x.view(-1, self.heads * self.out_channels)
+            x = self.convs[i](x, data.train_pos_edge_index).flatten(1)
             if i != self.num_layer - 2:
                 x = activate_func(x, self.args["act"])
                 # x = F.dropout(x, p=self.args["dropout"], training=self.training)
diff --git a/autogl/module/model/dgl/graphsage_dgl.py b/autogl/module/model/dgl/graphsage_dgl.py
index 6c18d71..433e0d7 100644
--- a/autogl/module/model/dgl/graphsage_dgl.py
+++ b/autogl/module/model/dgl/graphsage_dgl.py
@@ -1,6 +1,7 @@
 import torch
 import typing as _typing
 
+import torch.nn.functional as F
 from dgl.nn.pytorch.conv import SAGEConv
 import torch.nn.functional
 import autogl.data
@@ -48,11 +49,10 @@ class GraphSAGE(ClassificationSupportedSequentialModel):
             else:
                 self._dropout: _typing.Optional[torch.nn.Dropout] = None
 
-        def forward(self, data, enable_activation: bool = True) -> torch.Tensor:
-            x: torch.Tensor = data.ndata['x']
-            
+        def forward(self, data, x, enable_activation: bool = True) -> torch.Tensor:
+            # x = data.ndata['x']
             x: torch.Tensor = self._convolution.forward(data, x)
-            if self._activation_name is not None and enable_activation:
+            if (self._activation_name is not None) and enable_activation:
                 x: torch.Tensor = activate_func(x, self._activation_name)
             if self._dropout is not None:
                 x: torch.Tensor = self._dropout.forward(x)
@@ -142,7 +142,7 @@ class GraphSAGE(ClassificationSupportedSequentialModel):
                             hidden_features[i],
                             num_classes,
                             aggr,
-                            _layers_dropout[i + 1],
+                            dropout_probability=_layers_dropout[i + 1],
                         )
                     )
 
@@ -197,6 +197,15 @@ class GraphSAGE(ClassificationSupportedSequentialModel):
     def lp_decode_all(self, z):
         prob_adj = z @ z.t()
         return (prob_adj > 0).nonzero(as_tuple=False).t()
+    
+    def forward(self, data):
+        # only for test 
+        x = data.ndata['x']
+        for i in range(len(self.__sequential_encoding_layers)):
+            x = self.__sequential_encoding_layers[i](data,x)
+
+        return F.log_softmax(x, dim=1)
+
 
 
 @register_model("sage")
diff --git a/test/model_nlf/nclf_dgl.py b/test/model_nlf/nclf_dgl.py
index d1bdcb6..1bd39db 100644
--- a/test/model_nlf/nclf_dgl.py
+++ b/test/model_nlf/nclf_dgl.py
@@ -6,7 +6,6 @@ from tqdm import tqdm
 import time
 
 sys.path.append("../../")
-print(os.getcwd())
 os.environ["AUTOGL_BACKEND"] = "dgl"
 # os.environ["AUTOGL_BACKEND"] = "pyg"
 from autogl.backend import DependentBackend
@@ -17,13 +16,13 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
 
-from autogl.module.model import GCN
+from autogl.module.model import GAT,GraphSAGE
 
 from pdb import set_trace
 import numpy as np
 from autogl.solver.utils import set_seed
 set_seed(202106)
-
+import argparse
 
 def evaluate(model, graph, labels, mask):
     model.eval()
@@ -37,6 +36,7 @@ def evaluate(model, graph, labels, mask):
 
 
 def main():
+    
 
     # set up seeds, args.seed supported
     torch.manual_seed(seed=202106)
@@ -59,12 +59,23 @@ def main():
     labels = data.ndata['label']
     n_edges = data.number_of_edges()
 
-    model = GCN(data.ndata['x'].size(1), dataset.num_classes, [16], activation_name='relu',
-                dropout = 0.5).to(device)
+    args={}
+    args["features_num"]=data.ndata['x'].size(1)
+    args['hidden']=[16]
+    args["heads"]=8
+    args['dropout']=0.6
+    args["num_class"]=dataset.num_classes
+    args["num_layers"]=2
+    args['act']='relu'
+
+
+    # model = GAT(args)
+    model = GraphSAGE(args["features_num"],
+                      args["num_class"],
+                      [16],'relu',0.5)
 
     criterion = nn.CrossEntropyLoss()  # defaul reduce is true
     optimizer = optim.Adam(model.parameters(), lr=0.01)
-    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
 
     dur = []
     for epoch in range(200):