Fix infer

heyufan1995 · heyufan1995 · commit 02f7388923c0 · 2024-07-08T10:38:14.000-04:00
Signed-off-by: heyufan1995 &lt;heyufan1995@gmail.com&gt;
diff --git a/scripts/infer.py b/scripts/infer.py
@@ -168,6 +168,7 @@ def infer(
             batch_data = self.batch_data
         else:
             batch_data = self.infer_transforms(image_file)
+            batch_data["label_prompt"] = label_prompt
             batch_data = list_data_collate([batch_data])
             self.batch_data = batch_data
         if point is not None:
