Bugfix sybil-predict

Properly set --return-attentions when --write-attention-images is set.
reginabarzilaygroup · Jul 30, 2024 · 81a412c · 81a412c
1 parent 0d85c12
commit 81a412c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sybil/predict.py b/sybil/predict.py
@@ -49,7 +49,6 @@ def _get_parser():
         help="Generate images with attention overlap. Sets --return-attentions (if not already set).",
     )
 
-
     parser.add_argument(
         "--file-type",
         default="auto",
@@ -90,6 +89,8 @@ def predict(
 ):
     logger = sybil.utils.logging_utils.get_logger()
 
+    return_attentions |= write_attention_images
+
     input_files = os.listdir(image_dir)
     input_files = [os.path.join(image_dir, x) for x in input_files if not x.startswith(".")]
     input_files = [x for x in input_files if os.path.isfile(x)]