[Doc] Fix the doc for the partition_graph function (#4708)

* Fix the doc for the partition_graph func

* Lint for a long line

* Fix a lint
Co-authored-by: Triston Cao <tristonc@nvidia.com>

python/dgl/distributed/partition.py

@@ -569,8 +569,9 @@ def partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method=
     ...                                 out_path='output/', reshuffle=True,
     ...                                 balance_ntypes=g.ndata['train_mask'],
     ...                                 balance_edges=True)
-    >>> g, node_feats, edge_feats, gpb, graph_name = dgl.distributed.load_partition(
-    ...                                 'output/test.json', 0)
+    >>> (
+    ...     g, node_feats, edge_feats, gpb, graph_name, ntypes_list, etypes_list,
+    ... ) = dgl.distributed.load_partition('output/test.json', 0)
     '''
     def get_homogeneous(g, balance_ntypes):
         if g.is_homogeneous: