I need to experiment with a few options, but the main function will be to expand the size of the segmented region identified by the network (which will probably use skimage.draw, amongst other functions)
This is my most recent attempt:
def modify_function(y_pred):

    # Convert from one-hot encoding
    y_pred_modified = np.argmax(y_pred,axis=-1)
    
    # Do some further processing in numpy / scipy / skimage...
    
    return y_pred_modified

class TubeNetModel(keras.Model):

    def train_step(self, data):

        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass. Result is one-hot encoded with two categories (0 and 1)

            # Call the function to modify the output of the network
            y_pred_modified = tf.numpy_function(modify_function,[y_pred],y_pred.dtype)
            # Convert result back to one-hot encoding
            y_pred_modified = tf.one_hot(tf.cast(y_pred_modified, tf.int32), 2)
            # Reshape (necessary?)
            y_pred_modified = tf.reshape(y_pred_modified,tf.shape(y_pred))
             
            # Calculate the loss   
            loss = self.compiled_loss(y, y_pred_modified, regularization_losses=self.losses)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        # Update metrics (includes the metric that tracks the loss)
        self.compiled_metrics.update_state(y, y_pred)
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

However, this produces the following error when the apply_gradients method is called:
ValueError: No gradients provided for any variable: 

followed by a long list of model layer instances. I think this is due to the numpy function call failing and therefore no gradeints being calculated, but I’m not even sure how to debug it! Any suggestions?