diff --git a/tensorflow_gnn/graph/graph_piece.py b/tensorflow_gnn/graph/graph_piece.py
index a158be0f..e4709870 100644
--- a/tensorflow_gnn/graph/graph_piece.py
+++ b/tensorflow_gnn/graph/graph_piece.py
@@ -281,7 +281,7 @@ def update_fn(value):
                                       new_spec._metadata)
       return value
 
-    return self.__class__(
+    return self.__class__(  # pytype: disable=not-instantiable
         tf.nest.map_structure(update_fn, self._data), new_spec)
 
   @property