C4_W4: Neural_Style_Transfer train_step

Art_Generation_with_Neural_Style_Transfer, Exercise 6 - train_step: Can someone help me understand this function? I’m still very confused although I was able to get full grade.

  1. how is generated_image updated? I don’t see it got assigned anywhere inside the function.
  2. why do we need to use @tf.function()?
  3. Why cannot we use the old way of training the model like calling model.fit like how we did in the unet model?

So what we did there was build a custom training loop for transferring the style.

  1. We’re updating the image in the line where we’re applying the gradients to it later in the function
  2. We could do it without tf.function, but the execution being in eager mode will be slow. With tf. function, it’ll be faster because that converts the function into a static graph.
  3. You can actually use model.fit with custom training loops, but it’d have to be defined differently than here.

I’ll link a couple readings from tensorflow’s documentation for you to understand custom training loops better.


When we compute a_G as the vgg_model_outputs for the current generated image, we get a list of tensor.

I am not able to see how to use that within compute_style_cost and compute_layer_style_cost, as they take tensors as inputs.

Can anyone please guide me on this?

I did try to read the above links provided but was not able to see how it is useful in our case.

Thanks in advance.

If you read the code in compute_style_cost(), you will find that it loops through the items in the list.