fixing retry script

pyg-team · Feb 3, 2025 · 335f20e · 335f20e
1 parent b9dc0fc
commit 335f20e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_geometric/nn/nlp/txt2kg.py b/torch_geometric/nn/nlp/txt2kg.py
@@ -164,20 +164,20 @@ def add_doc_2_KG(
                                           _chunk_to_triples_str_cloud,
                                           self.NVIDIA_API_KEY, self.NIM_MODEL),
                                     nprocs=num_procs)
+                                break
                             except:  # noqa
                                 # keep retrying, txt2kg is costly -> stoppage is costly
                                 pass
-                            break
 
                         # Collect the results from each process
                         self.relevant_triples[key] = []
                         for rank in range(num_procs):
                             self.relevant_triples[key] += torch.load(
                                 "/tmp/outs_for_proc_" + str(rank))
                             os.remove("/tmp/outs_for_proc_" + str(rank))
+                        break
                     except:
                         pass
-                    break
         # Increment the doc_id_counter for the next document
         self.doc_id_counter += 1