code: def dgi(generator, gen, model_type):
      tf.random.set_seed(1234)
      np.random.seed(1234)
  
      emb_dim = 4
  
      base_model = model_type(
          generator=generator,
          activations=["relu"],
          layer_sizes=[emb_dim],
      )
      infomax = DeepGraphInfomax(base_model)
  
      model = tf.keras.Model(*infomax.in_out_tensors())
      model.compile(
          loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam"
      )
      model.fit(gen)
      return model
  
file_path: stellargraph/tests/reproducibility/test_deep_graph_infomax.py
imports: from .fixtures import assert_reproducible
  from stellargraph.layer import DeepGraphInfomax, GCN, APPNP, GAT, PPNP
  from stellargraph.mapper import FullBatchNodeGenerator, CorruptedGenerator
  from ..test_utils.graphs import example_graph_random
  from ..None import require_gpu
  import tensorflow
  import pytest
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: dgi(generator, gen, model_type) and returns: model