diff --git a/deepwalk/__main__.py b/deepwalk/__main__.py
index 1070e07..149b796 100644
--- a/deepwalk/__main__.py
+++ b/deepwalk/__main__.py
@@ -72,7 +72,7 @@ def process(args):
     walks = graph.build_deepwalk_corpus(G, num_paths=args.number_walks,
                                         path_length=args.walk_length, alpha=0, rand=random.Random(args.seed))
     print("Training...")
-    model = Word2Vec(walks, size=args.representation_size, window=args.window_size, min_count=0, sg=1, hs=1, workers=args.workers)
+    model = Word2Vec(walks, vector_size=args.representation_size, window=args.window_size, min_count=0, sg=1, hs=1, workers=args.workers)
   else:
     print("Data size {} is larger than limit (max-memory-data-size: {}).  Dumping walks to disk.".format(data_size, args.max_memory_data_size))
     print("Walking...")
diff --git a/deepwalk/skipgram.py b/deepwalk/skipgram.py
index 42c0770..7350953 100644
--- a/deepwalk/skipgram.py
+++ b/deepwalk/skipgram.py
@@ -3,12 +3,11 @@
 import logging
 from multiprocessing import cpu_count
 from six import string_types
-
 from gensim.models import Word2Vec
-from gensim.models.word2vec import Vocab
 
 logger = logging.getLogger("deepwalk")
 
+
 class Skipgram(Word2Vec):
     """A subclass to allow more customization of the Word2Vec internals."""
 
@@ -24,7 +23,7 @@ def __init__(self, vocabulary_counts=None, **kwargs):
         kwargs["sg"] = 1
         kwargs["hs"] = 1
 
-        if vocabulary_counts != None:
-          self.vocabulary_counts = vocabulary_counts
+        if vocabulary_counts is not None:
+            self.vocabulary_counts = vocabulary_counts
 
         super(Skipgram, self).__init__(**kwargs)