fix: return loaded parameters in Aria load_weights (#57476)

Signed-off-by: Debanshu Das <debanshu@google.com>
2026-09-29 04:05:07 +00:00
parent 295a0b8d64
commit 13c8589b86
1 changed files with 5 additions and 2 deletions
@@ -526,6 +526,9 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
        logits = self.logits_processor(self.lm_head, hidden_states)
        return logits

-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
        loader = AutoWeightsLoader(self)
-        loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)