Only attempt download of completed tv2 models

GitOrigin-RevId: 5eee3b89c1926deb8a7e325a5a5203a435b50d1d
gretelai · Jan 22, 2024 · c832c4f · c832c4f
1 parent a8a3a08
commit c832c4f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/gretel_trainer/relational/multi_table.py b/src/gretel_trainer/relational/multi_table.py
@@ -602,8 +602,9 @@ def transform_v2(
 
         output_tables = {}
         for table, model in self._transforms_train.models.items():
-            with model.get_artifact_handle("data_preview") as data_preview:
-                output_tables[table] = pd.read_csv(data_preview)
+            if table in task.completed:
+                with model.get_artifact_handle("data_preview") as data_preview:
+                    output_tables[table] = pd.read_csv(data_preview)
 
         self._post_process_transformed_tables(
             output_tables=output_tables,