class_docstr: Tests for ops in horovod.tensorflow.keras.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow.keras","."]
code: def train_model(self, backward_passes_per_step):
      with self.test_session(config=self.config) as sess:
          K.set_session(sess)
  
          opt = keras.optimizers.RMSprop(lr=0.0001)
          opt = hvd.DistributedOptimizer(
              opt,
              backward_passes_per_step=backward_passes_per_step,
              average_aggregated_gradients=True,
          )
  
          model = keras.models.Sequential()
          model.add(keras.layers.Dense(2, input_shape=(3,)))
          model.add(keras.layers.RepeatVector(3))
          model.add(keras.layers.ThresholdedReLU(0.5))
          model.compile(
              loss=keras.losses.mean_squared_error,
              optimizer=opt,
              metrics=[keras.metrics.categorical_accuracy],
              sample_weight_mode="temporal",
          )
  
          x = np.random.random((1, 3))
          y = np.random.random((1, 3, 3))
  
          def generator():
              while 1:
                  yield (x, y)
  
          # No assertions, we just need to verify that it doesn't hang
          callbacks = [
              hvd.callbacks.BroadcastGlobalVariablesCallback(0)
          ]
          model.fit_generator(
              generator(),
              steps_per_epoch=10,
              callbacks=callbacks,
              epochs=0,
              verbose=0,
              workers=4,
              initial_epoch=1,
          )
  
imports: import math
  import numpy
  import os
  import pytest
  import sys
  import tensorflow
  import warnings
  from packaging import version
  from tensorflow import keras
  from horovod.common.util import is_version_greater_equal_than
  import horovod.tensorflow.keras
  from common import temppath
label_desc: Use image node_id 2 for calling the TfKerasTests obj's underlying member method code with example usage: obj.train_model(backward_passes_per_step) without return types