Parallelization support

raven_modeling_minimal.py

@@ -492,7 +492,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         attn_maps: dict = {},
         return_attn: bool = False,
     ):
-        x = self.transformer.adapter(torch.cat([x, input_embeds], dim=-1))
+        x = self.transformer.adapter(torch.cat([x, input_embeds.to(x.device)], dim=-1))
         for idx, block in enumerate(self.transformer.core_block, start=1):
             x, attn_map = block(x, freqs_cis, block_idx + idx, mask, past_key_values, return_attn=return_attn)
             attn_maps[block_idx + idx] = attn_map