fix nonetype in atkgen probe (#538)

NVIDIA · Mar 4, 2024 · a7a1e00 · a7a1e00
1 parent f1167c7
commit a7a1e00
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/garak/probes/atkgen.py b/garak/probes/atkgen.py
@@ -136,7 +136,13 @@ def probe(self, generator) -> List[garak.attempt.Attempt]:
                         f" turn {t.n:02}: waiting for [{generator.name[:10]:<10}]"
                     )
                 # send the challenge and get the response
-                response = generator.generate(challenge)[0].strip()
+                try:
+                    response = generator.generate(challenge)[0].strip()
+                except AttributeError as ae:
+                    if generator.generate(challenge)[0] is None:
+                        response = ""
+                    else:
+                        raise AttributeError from ae
                 # log the response
                 turn = ("model", response)
                 turns.append(turn)