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?