def __init__(self, initializer='ones'):
    super(SubclassedKerasModel, self).__init__()
    self._can_use_graph_functions = True
    self.layer_a = keras.layers.Dense(64, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_b = keras.layers.Dense(128, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_c = keras.layers.Dense(256, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_d = keras.layers.Dense(256, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_e = keras.layers.Dense(10, kernel_initializer=initializer, bias_initializer='zeros')
def __init__(self, initializer='ones'):
    super(SubclassedKerasModel, self).__init__()
    self._can_use_graph_functions = True
    self.layer_a = keras.layers.Dense(64, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_b = keras.layers.Dense(128, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_c = keras.layers.Dense(256, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_d = keras.layers.Dense(256, kernel_initializer=initializer, bias_initializer='zeros')
    self.layer_e = keras.layers.Dense(10, kernel_initializer=initializer, bias_initializer='zeros')