diff --git a/lab2/Part2_FaceDetection.ipynb b/lab2/Part2_FaceDetection.ipynb
index c877149c..59678b73 100644
--- a/lab2/Part2_FaceDetection.ipynb
+++ b/lab2/Part2_FaceDetection.ipynb
@@ -779,7 +779,7 @@
         "    y_logit, z_mean, z_logsigma, x_recon = ss_vae(x)\n",
         "\n",
         "    '''TODO: call the SS_VAE loss function to compute the loss'''\n",
-        "    loss, class_loss = ss_vae_loss_function('''TODO arguments''') # TODO\n",
+        "    loss, class_loss, _ = ss_vae_loss_function('''TODO arguments''') # TODO\n",
         "  \n",
         "  '''TODO: use the GradientTape.gradient method to compute the gradients.\n",
         "     Hint: this is with respect to the trainable_variables of the SS_VAE.'''\n",