diff --git a/pygcn/models.py b/pygcn/models.py
index 91b2fcf..f701029 100644
--- a/pygcn/models.py
+++ b/pygcn/models.py
@@ -15,4 +15,4 @@ def forward(self, x, adj):
         x = F.relu(self.gc1(x, adj))
         x = F.dropout(x, self.dropout, training=self.training)
         x = self.gc2(x, adj)
-        return F.log_softmax(x)
+        return F.log_softmax(x, dim =1)