Upload folder using huggingface_hub

modeling_tiny_mixtral.py

@@ -179,7 +179,6 @@ def apply_rotary_embeddings(x:torch.Tensor,freq_complex:torch.Tensor,device:str)
     Returns:
         torch.Tensor: The tensor after applying Rotary Position Embeddings.
     """
-    freq_complex = freq_complex.to(x.device)
     x_complex=torch.view_as_complex(x.float().reshape(*x.shape[:-1],-1,2)) #N,seq_len,h,head_dim/2,2 
     
     freq_complex=freq_complex.unsqueeze(0).unsqueeze(2) # 1,seq_len,1,head_dim/2
@@ -188,7 +187,7 @@ def apply_rotary_embeddings(x:torch.Tensor,freq_complex:torch.Tensor,device:str)
     x_out=torch.view_as_real(x_rotated) #(N,seq_len,h,head_dim/2,2)
     x_out=x_out.reshape(*x.shape)
     
-    return x_out.type_as(x).to(x.device)
+    return x_out.type_as(x).to(device)
 
 
 
@@ -307,10 +306,6 @@ class SimpleMultiHeadAttention(nn.Module):
             
             q = q_rotary.transpose(1, 2)  # Back to (batch_size, num_heads, seq_len, head_dim)
             k = k_rotary.transpose(1, 2)  # Back to (batch_size, num_heads, seq_len, head_dim)
-            
-            q = q.to(self.device)
-            k = k.to(self.device)
-            v = v.to(self.device)
         
         # Compute attention
         if self.flash:
@@ -340,7 +335,6 @@ class SimpleMultiHeadAttention(nn.Module):
         
         # Reshape back to (batch_size, seq_len, dim)
         y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, self.dim)
-        y = y.to(self.device)
         
         # Output projection
         y = self.resid_dropout(self.c_proj(y))
@@ -519,7 +513,7 @@ class layer(nn.Module):
         out=h+ffn_output
         
         
-        return out
+        return out, router_loss
         
 
 class tiny_mixtral(nn.Module):
@@ -584,7 +578,6 @@ class TinyMixtralForCausalLM(PreTrainedModel, GenerationMixin):
         self.config = config
         self.post_init()
 
-
     def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
 
         outputs, load_balancing_loss = self.model(input_ids, start_pos=0)