Upload ModularStarEncoder

modularStarEncoder.py

@@ -1,6 +1,6 @@
 from transformers import  Starcoder2Model
 import sys
-from .config import ModularStarEncoderConfig
+from config import ModularStarEncoderConfig
 import os
 from dataclasses import dataclass
 from typing import Optional, Tuple, Union
@@ -171,8 +171,11 @@ class StarEncoder2PreTrainingHeads(nn.Module):
 
     def forward(self, sequence_output, pooled_output,idx_layer: Optional[torch.Tensor] = None):
         if self.is_matryoshka:
-            prediction_scores = self.predictions(torch.cat([sequence_output , self.conditional_embeddings(torch.tensor(idx_layer,device=sequence_output.get_device()).int()).expand(sequence_output.size()[0],sequence_output.size()[1],-1)],dim=-1))
-            seq_relationship_score = self.seq_relationship(torch.cat([pooled_output , self.conditional_embeddings(torch.tensor(idx_layer,device=pooled_output.get_device()).int()).expand(pooled_output.size()[0],-1)],dim=-1))
+            device_sequence = sequence_output.get_device()
+            if device_sequence<0:
+                device_sequence = "cpu"
+            prediction_scores = self.predictions(torch.cat([sequence_output , self.conditional_embeddings(torch.tensor(idx_layer,device=device_sequence).int()).expand(sequence_output.size()[0],sequence_output.size()[1],-1)],dim=-1))
+            seq_relationship_score = self.seq_relationship(torch.cat([pooled_output , self.conditional_embeddings(torch.tensor(idx_layer,device=device_sequence).int()).expand(pooled_output.size()[0],-1)],dim=-1))
         else:
             prediction_scores = self.predictions(sequence_output)
             seq_relationship_score = self.seq_relationship(pooled_output)