fix: partition adapter mask when batch size is specified (#41)

- fix: partition adapter mask when batch size is specified (e6e3a6f1dcfa01de524099f84cd7a85299175f55)

modeling_xlm_roberta.py

@@ -558,15 +558,18 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             )
         else:
             range_iter = range(0, len(sentences), batch_size)
-        lora_arguments = (
-            {"adapter_mask": adapter_mask} if adapter_mask is not None else {}
-        )
+
         for i in range_iter:
             encoded_input = self.tokenizer(
                 sentences[i : i + batch_size],
                 return_tensors="pt",
                 **tokenizer_kwargs,
             ).to(self.device)
+            lora_arguments = (
+                {"adapter_mask": adapter_mask[i : i + batch_size]}
+                if adapter_mask is not None
+                else {}
+            )
             token_embs = self.forward(**encoded_input, **lora_arguments)[0]
 
             # Accumulate in fp32 to avoid overflow