TypeError: Can not convert a NoneType into a Tensor or Operation -- Error believe related to converting to graph

Question

Below find my model:

class CustomModel(tf.keras.Model):
    def __init__(self, model1, model2, model3, model4):
        super(deep_and_wide, self).__init__()
        self.model1 = model1
        self.model2 = model2
        self.model3 = model3
        self.model4 = model4

    def call(self, inputs):
        x1 = self.mode1([inputs["a"], inputs["b"]])
        x2 = self.model2([inputs["a"], inputs["b"]])
        x3 = self.model3([inputs["a"], inputs["b"]])
        x4 = self.model4([inputs["a"], inputs["b"]])
        x = Concatenate()([x1, x2, x3])
        x = TimeDistributed(Dense(2))(x)

        x = Add()([x, x4])
        x_fc = Dense(1)(x)
        x_ec = Dense(1)(x)

        return x_fc, x_ec


    def train_step(self, data):
        with tf.GradientTape() as tape:
            data = data_adapter.expand_1d(data)
            batch_inputs, batch_outputs, sample_weight= data_adapter.unpack_x_y_sample_weight(data)

            y_true_fc, y_true_ec = batch_outputs["y_fc"], batch_outputs["y_ec"]
            y_pred_fc, y_pred_ec = self(batch_inputs, training=True)
            loss_fc = self.compiled_loss(y_true_fc, y_pred_fc)
            loss_ec = self.compiled_loss(y_true_ec, y_pred_ec)
        print("here")
        trainable_variables = self.trainable_variables
        print("here")
        gradients = tape.gradient([loss_fc, loss_ec], trainable_variables)
        print("here")
        self.optimizer.apply_gradients(zip(gradients, trainable_variables))
        print("here")

And below is my custom loss

class CustomLoss(tf.keras.losses.Loss):
    def __init__(self, mask=True, alpha=1, beta=1, gamma=1, dtype=tf.float64):
        super(CustomLoss, self).__init__(reduction=tf.keras.losses.Reduction.NONE)
        self.mask = mask
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.dtype = dtype

    def call(self, y_true, y_pred):
        def loss_fn(y_true, y_pred, mask):
            y_true = tf.boolean_mask(y_true, mask)
            y_pred = tf.boolean_mask(y_pred, mask)
            return tf.keras.losses.MSE(y_true, y_pred)

        self.mask = tf.not_equal(y_true, 0.)

        y_true = tf.cast(y_true, self.dtype)
        y_pred = tf.cast(y_pred, self.dtype)

        y_pred = tf.multiply(y_pred, tf.cast(self.mask, dtype=self.dtype))
        y_pred_cum = tf.math.cumsum(y_pred, axis=1)
        y_pred_cum = tf.multiply(y_pred_cum, tf.cast(self.mask, dtype=self.dtype))

        y_true_cum = tf.math.cumsum(y_true, axis=1)
        y_true_cum = tf.multiply(y_true_cum, tf.cast(self.mask, dtype=self.dtype))

        loss_value = self.alpha * loss_fn(y_true, y_pred, self.mask) + \
                     self.gamma * loss_fn(y_true_cum, y_pred_cum, self.mask)

        return loss_value

And then finally:

optimizer = tf.keras.optimizers.Adam()

loss = CustomLoss()
model.compile(optimizer, loss)

model.fit(train_data, epochs=5, validation_data=val_data)

My data inputs are of size (sequence length, feature length) where sequence length is variable hence I am using tf.data.experimental.bucket_by_sequence_length to pad to max sequence length of the batch (as opposed to batch to max sequence length). All in all, my train and val data are tf.data.Datasets each created using tf.data.experimental.bucket_by_sequence_length where each batch is of size (None, None, feature length).

When I run the above code, I get the following errors and cannot seem to understand where I am going wrong:

Traceback (most recent call last):
  File "<input>", line 75, in <module>
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\keras\engine\training.py", line 1100, in fit
    tmp_logs = self.train_function(iterator)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\def_function.py", line 725, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\eager\def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
TypeError: in user code:
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\keras\engine\training.py:805 train_function  *
        return step_function(self, iterator)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\keras\engine\training.py:795 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\distribute\distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\distribute\distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\distribute\distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\keras\engine\training.py:790 run_step  **
        with ops.control_dependencies(_minimum_control_deps(outputs)):
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\ops.py:5359 control_dependencies
        return get_default_graph().control_dependencies(control_inputs)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\func_graph.py:362 control_dependencies
        return super(FuncGraph, self).control_dependencies(filtered_control_inputs)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\ops.py:4815 control_dependencies
        c = self.as_graph_element(c)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\ops.py:3726 as_graph_element
        return self._as_graph_element_locked(obj, allow_tensor, allow_operation)
    C:\Users\\Anaconda3\envs\tf_recsys\lib\site-packages\tensorflow\python\framework\ops.py:3814 _as_graph_element_locked
        raise TypeError("Can not convert a %s into a %s." %
    TypeError: Can not convert a NoneType into a Tensor or Operation.

The four print statements inserted in the train_step function above are printed.

Answer 1

This NoneType refers to the returned value of the custom train_step, when using a custom train_step you should return something that can be converted into a tensor so that the minimum control dependencies can process it, typically, the loss value as {"loss": loss_value} and potentially some other metrics, or at least an empty dict {}.