def METHOD_NAME(self):
    K.clear_session()
    state_dim = (2, 3)
    action_dim = 5
    inp = Input(shape=(state_dim))
    # Determine whether we are dealing with an image input (Atari) or not
    if (len(state_dim) > 2):
        inp = Input((state_dim[1:]))
        x = conv_block(inp, 32, (2, 2), 8)
        x = conv_block(x, 64, (2, 2), 4)
        x = conv_block(x, 64, (2, 2), 3)
        x = Flatten()(x)
        x = Dense(256, activation='relu')(x)
    else:
        x = Flatten()(inp)
        x = Dense(64, activation='relu')(x)
        x = Dense(64, activation='relu')(x)
    x = Dense(action_dim + 1, activation='linear')(x)
    x = Lambda(lambda i: K.expand_dims(i[:,0],-1) + i[:,1:] - K.mean(i[:,1:], keepdims=True),
               output_shape=(action_dim,))(x)
    keras_model = Model(inp, x)
    data = np.random.rand(1000, 2, 3).astype(np.float32)
    expected = keras_model.predict(data)
    onnx_model = mock_keras2onnx.convert_keras(keras_model, keras_model.name)
    self.assertTrue(
        run_keras_and_ort(onnx_model.graph.name, onnx_model, keras_model, data, expected,
                          self.model_files, rtol=1e-2, atol=2e-2))