diff --git a/python/graphstorm/gsf.py b/python/graphstorm/gsf.py
index 2758875538..c1757f4655 100644
--- a/python/graphstorm/gsf.py
+++ b/python/graphstorm/gsf.py
@@ -584,7 +584,8 @@ def set_encoder(model, g, config, train_task):
                                  num_ffn_layers_in_gnn=config.num_ffn_layers_in_gnn)
     elif model_encoder_type == "sage":
         # we need to check if the graph is homogeneous
-        assert check_homo(g) == True, 'The graph is not a homogeneous graph'
+        assert check_homo(g) == True, \
+            'The graph is not a homogeneous graph, can not use sage model encoder'
         # we need to set the num_layers -1 because there is an output layer that is hard coded.
         gnn_encoder = SAGEEncoder(h_dim=config.hidden_size,
                                   out_dim=config.hidden_size,
@@ -594,6 +595,9 @@ def set_encoder(model, g, config, train_task):
                                   num_ffn_layers_in_gnn=config.num_ffn_layers_in_gnn,
                                   norm=config.gnn_norm)
     elif model_encoder_type == "gat":
+        # we need to check if the graph is homogeneous
+        assert check_homo(g) == True, \
+            'The graph is not a homogeneous graph, can not use gat model encoder'
         gnn_encoder = GATEncoder(h_dim=config.hidden_size,
                                  out_dim=config.hidden_size,
                                  num_heads=config.num_heads,
@@ -601,6 +605,9 @@ def set_encoder(model, g, config, train_task):
                                  dropout=dropout,
                                  num_ffn_layers_in_gnn=config.num_ffn_layers_in_gnn)
     elif model_encoder_type == "gatv2":
+        # we need to check if the graph is homogeneous
+        assert check_homo(g) == True, \
+            'The graph is not a homogeneous graph, can not use gatv2 model encoder'
         gnn_encoder = GATv2Encoder(h_dim=config.hidden_size,
                                    out_dim=config.hidden_size,
                                    num_heads=config.num_heads,
diff --git a/python/graphstorm/model/gat_encoder.py b/python/graphstorm/model/gat_encoder.py
index 3a590a294c..7b1134d7df 100644
--- a/python/graphstorm/model/gat_encoder.py
+++ b/python/graphstorm/model/gat_encoder.py
@@ -162,7 +162,7 @@ class GATEncoder(GraphConvEncoder):
         Hidden dimension
     out_dim : int
         Output dimension
-    num_head : int
+    num_heads : int
         Number of multi-heads attention
     num_hidden_layers : int
         Number of hidden layers. Total GNN layers is equal to num_hidden_layers + 1. Default 1
@@ -170,6 +170,8 @@ class GATEncoder(GraphConvEncoder):
         Dropout. Default 0.
     activation : callable, optional
         Activation function. Default: None
+    last_layer_act: bool
+        Whether call activation function in the last GNN layer.
     num_ffn_layers_in_gnn: int
         Number of ngnn gnn layers between GNN layers
     """
diff --git a/python/graphstorm/model/gatv2_encoder.py b/python/graphstorm/model/gatv2_encoder.py
index 2ea7eda4e7..a6b8a6b665 100644
--- a/python/graphstorm/model/gatv2_encoder.py
+++ b/python/graphstorm/model/gatv2_encoder.py
@@ -166,7 +166,7 @@ class GATv2Encoder(GraphConvEncoder):
         Hidden dimension
     out_dim : int
         Output dimension
-    num_head : int
+    num_heads : int
         Number of multi-heads attention
     num_hidden_layers : int
         Number of hidden layers. Total GNN layers is equal to num_hidden_layers + 1. Default 1
@@ -174,6 +174,8 @@ class GATv2Encoder(GraphConvEncoder):
         Dropout. Default 0.
     activation : callable, optional
         Activation function. Default: None
+    last_layer_act: bool
+        Whether call activation function in the last GNN layer.
     num_ffn_layers_in_gnn: int
         Number of ngnn gnn layers between GNN layers
     """