class_docstr: Tests for ops in horovod.tensorflow.keras.
class_docstr_tok: ["Tests","for","ops","in","horovod.tensorflow.keras","."]
code: def test_load_model_broadcast(self):
      def create_model():
          opt = keras.optimizers.SGD(lr=0.01 * hvd.size(), momentum=0.9)
          opt = hvd.DistributedOptimizer(opt)
  
          model = keras.models.Sequential()
          model.add(keras.layers.Dense(2, input_shape=(3,)))
          model.add(keras.layers.RepeatVector(3))
          model.add(keras.layers.TimeDistributed(keras.layers.Dense(3)))
          model.compile(
              loss=keras.losses.MSE,
              optimizer=opt,
              metrics=[keras.metrics.categorical_accuracy],
              sample_weight_mode="temporal",
          )
  
          return model
  
      with temppath() as fname:
          with self.session(config=self.config) as sess:
              K.set_session(sess)
  
              model = create_model()
  
              x = np.random.random((1, 3))
              y = np.random.random((1, 3, 3))
              model.train_on_batch(x, y)
  
              if hvd.rank() == 0:
                  model.save(fname)
  
          K.clear_session()
          with self.session(config=self.config) as sess:
              K.set_session(sess)
  
              weight = np.random.random((1, 3))
  
              if hvd.rank() == 0:
                  model = hvd.load_model(fname)
              else:
                  model = create_model()
  
              def generator():
                  while 1:
                      yield (x, y, weight)
  
              if hvd.rank() == 0:
                  self.assertEqual(len(model.optimizer.weights), 5)
              else:
                  self.assertEqual(len(model.optimizer.weights), 0)
  
              # No assertions, we just need to verify that it doesn't hang
              callbacks = [
                  hvd.callbacks.BroadcastGlobalVariablesCallback(0)
              ]
              model.fit_generator(
                  generator(),
                  steps_per_epoch=1,
                  callbacks=callbacks,
                  epochs=1,
                  verbose=0,
                  workers=4,
                  initial_epoch=0,
              )
  
              self.assertEqual(len(model.optimizer.weights), 5)
  
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 9 for calling the TfKerasTests obj's underlying member method code with example usage: obj.test_load_model_broadcast() and returns: model