fix: distribute properly models

experimaestro · Dec 7, 2023 · fbc9b5c · fbc9b5c
1 parent 2c0e23d
commit fbc9b5c
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/src/xpmir/neural/cross.py b/src/xpmir/neural/cross.py
@@ -14,6 +14,9 @@
     DuoTwoStageRetriever,
     Retriever,
 )
+from xpmir.utils.utils import easylog
+
+logger = easylog()
 
 
 class CrossScorer(LearnableScorer, DistributableModel):
@@ -51,7 +54,13 @@ def forward(self, inputs: BaseRecords, info: TrainerContext = None):
         return self.classifier(pairs).squeeze(1)
 
     def distribute_models(self, update):
-        self.encoder.model = update(self.encoder.model)
+        if isinstance(self.encoder, DistributableModel):
+            self.encoder = self.distribute_models(self.encoder, update)
+        else:
+            logger.warning(
+                "Cross-encoder encoder is not distributable: "
+                "keeping it on one device"
+            )
 
 
 class DuoCrossScorer(DuoLearnableScorer, DistributableModel):