diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py
index 5677b37..d7c213c 100644
--- a/laborious/workflows/drift.py
+++ b/laborious/workflows/drift.py
@@ -32,6 +32,9 @@ class Drift:
             }
         }
 
+        model_config = input_data['model_config']
+        target_name = model_config['target']
+
         gathering_query = f"""
             SELECT *
             FROM {input_data['schema']}.{input_data['source_table_name']}
@@ -76,7 +79,7 @@ class Drift:
                 'reference_data': reference_data,
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
-                'target_name': input_data['target_name'],
+                'target_name': target_name,
                 'drift_metrics': input_data['drift_metrics'],
                 'chunk_period': input_data.get('chunk_period', 'min'),
             },
diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py
index 09411aa..82c9eb3 100644
--- a/tests/laborious/workflows/test_drift.py
+++ b/tests/laborious/workflows/test_drift.py
@@ -34,7 +34,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
         'source_table_name': 'test_source_table',
         'target_table_name': 'test_target_table',
         'interval': 60,
-        'target_name': 'test_target',
+        'model_config': {'target': 'test_target'},
         'drift_metrics': ['psi', 'ks'],
         'chunk_period': 'hour',
     }