diff --git a/template/{{.input_root_dir}}/{{template `project_name_alphanumeric_underscore` .}}/feature_engineering/notebooks/GenerateAndWriteFeatures.py.tmpl b/template/{{.input_root_dir}}/{{template `project_name_alphanumeric_underscore` .}}/feature_engineering/notebooks/GenerateAndWriteFeatures.py.tmpl
index da3f22d9..c8ac612c 100644
--- a/template/{{.input_root_dir}}/{{template `project_name_alphanumeric_underscore` .}}/feature_engineering/notebooks/GenerateAndWriteFeatures.py.tmpl	
+++ b/template/{{.input_root_dir}}/{{template `project_name_alphanumeric_underscore` .}}/feature_engineering/notebooks/GenerateAndWriteFeatures.py.tmpl	
@@ -82,8 +82,8 @@ pk_columns = dbutils.widgets.get("primary_keys")
 assert input_table_path != "", "input_table_path notebook parameter must be specified"
 assert output_table_name != "", "output_table_name notebook parameter must be specified"
 
-# Extract database name. Needs to be updated for Unity Catalog.
-output_database = output_table_name.split(".")[0]
+# Extract database name. Needs to be updated for Unity Catalog to the Schema name.
+output_database = output_table_name.split(".")[1]
 
 # COMMAND ----------