Changes for Core ML conversion

modelling_RW.py

@@ -29,7 +29,7 @@ logger = logging.get_logger(__name__)
 # In order not to degrade the quality of our HF-port, we keep these characteristics in the final model.
 class Linear(nn.Linear):
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        ret = input @ self.weight.T
+        ret = input @ self.weight.permute(1, 0)     #transpose(0, 1)  #.T
         if self.bias is None:
             return ret
         else:
@@ -68,7 +68,7 @@ class RotaryEmbedding(torch.nn.Module):
         self,
         seq_len: int,
         device="cuda",
-        dtype=torch.bfloat16,
+        dtype=torch.float16,
     ) -> torch.Tensor:
         if seq_len != self.seq_len_cached:
             self.seq_len_cached = seq_len
@@ -89,7 +89,7 @@ class RotaryEmbedding(torch.nn.Module):
 
     def forward(self, q, k):
         batch, seq_len, head_dim = q.shape
-        cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
+        cos, sin = self.cos_sin(seq_len, q.device)
         return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)