Merge branch 'latest-txt2kg' into rerank

pyg-team · Feb 5, 2025 · ae06392 · ae06392
2 parents cd34dcb + fb959cb
commit ae06392
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torch_geometric/loader/rag_loader.py b/torch_geometric/loader/rag_loader.py
@@ -109,10 +109,11 @@ def __init__(
                 (Default: obj:`''`)
         """
         fstore, gstore = data
-        assert len(raw_docs) == len(
-            embedded_docs), "Need raw and embedded docs to match"
         self.raw_docs = raw_docs
         self.k_for_docs = k_for_docs
+        if self.raw_docs:
+            assert len(raw_docs) == len(
+                embedded_docs), "Need raw and embedded docs to match"
         self.embedded_docs = embedded_docs
         self.feature_store = fstore
         self.graph_store = gstore