diff --git a/graph/src/main/scala/org/apache/spark/graph/GraphKryoRegistrator.scala b/graph/src/main/scala/org/apache/spark/graph/GraphKryoRegistrator.scala
index 13a22f9051..297506e1e3 100644
--- a/graph/src/main/scala/org/apache/spark/graph/GraphKryoRegistrator.scala
+++ b/graph/src/main/scala/org/apache/spark/graph/GraphKryoRegistrator.scala
@@ -2,23 +2,19 @@ package org.apache.spark.graph
 
 import com.esotericsoftware.kryo.Kryo
 
+import org.apache.spark.graph.impl.MessageToPartition
 import org.apache.spark.serializer.KryoRegistrator
 
 
 class GraphKryoRegistrator extends KryoRegistrator {
 
   def registerClasses(kryo: Kryo) {
-    //kryo.register(classOf[(Int, Float, Float)])
-    registerClass[Int, Int, Int](kryo)
+    kryo.register(classOf[Vertex[Object]])
+    kryo.register(classOf[Edge[Object]])
+    kryo.register(classOf[MutableTuple2[Object, Object]])
+    kryo.register(classOf[MessageToPartition[Object]])
 
     // This avoids a large number of hash table lookups.
     kryo.setReferences(false)
   }
-
-  private def registerClass[VD: Manifest, ED: Manifest, VD2: Manifest](kryo: Kryo) {
-    kryo.register(classOf[Vertex[VD]])
-    kryo.register(classOf[Edge[ED]])
-    kryo.register(classOf[MutableTuple2[VD, VD2]])
-    kryo.register(classOf[(Vid, VD2)])
-  }
 }