code: def test_WatchYourStep(barbell, weighted):
      generator = AdjacencyPowerGenerator(
          barbell, num_powers=5, weighted=weighted
      )
      gen = generator.flow(batch_size=4)
      wys = WatchYourStep(generator)
  
      x_in, x_out = wys.in_out_tensors()
  
      model = Model(inputs=x_in, outputs=x_out)
      model.compile(optimizer="adam", loss=graph_log_likelihood)
      model.fit(
          gen, epochs=1, steps_per_epoch=int(len(barbell.nodes()) // 4)
      )
  
      embs = wys.embeddings()
      assert embs.shape == (
          len(barbell.nodes()),
          wys.embedding_dimension,
      )
  
      # build() should always return tensors backed by the same trainable weights, and thus give the
      # same predictions
      preds1 = model.predict(gen, steps=8)
      preds2 = Model(*wys.in_out_tensors()).predict(gen, steps=8)
      np.testing.assert_array_equal(preds1, preds2)
  
file_path: stellargraph/tests/layer/test_watch_your_step.py
imports: from stellargraph.layer import AttentiveWalk, WatchYourStep
  import numpy
  from ..test_utils.graphs import barbell
  from stellargraph.mapper import AdjacencyPowerGenerator
  from stellargraph.losses import graph_log_likelihood
  import pytest
  from tensorflow.keras import Model
  from ..None import test_utils
label_desc: Use image node_id 5 for calling a global function with example usage: test_WatchYourStep(barbell, weighted) without return types