diff --git a/Train_Colab.ipynb b/Train_Colab.ipynb
index 553cba1..86e2f5d 100644
--- a/Train_Colab.ipynb
+++ b/Train_Colab.ipynb
@@ -192,6 +192,7 @@
         "\n",
         "img_size = 512\n",
         "upscale_attention = False\n",
+        "prediction_type = \"epsilon\"\n",
         "if model_type == \"SD1x\":\n",
         "    inference_yaml = \"v1-inference.yaml\"\n",
         "elif model_type == \"SD2_512_base\":\n",
@@ -199,6 +200,7 @@
         "    inference_yaml = \"v2-inference.yaml\"\n",
         "elif model_type == \"SD21\":\n",
         "    upscale_attention = True\n",
+        "    prediction_type = \"v_prediction\"\n",
         "    inference_yaml = \"v2-inference-v.yaml\"\n",
         "    img_size = 768\n",
         "\n",
@@ -209,7 +211,7 @@
         "--original_config_file {inference_yaml} \\\n",
         "--image_size {img_size} \\\n",
         "--checkpoint_path {base_path} \\\n",
-        "--prediction_type epsilon \\\n",
+        "--prediction_type {prediction_type} \\\n",
         "--upcast_attn False \\\n",
         "--dump_path {save_name}"
       ]