Upload model.py

code/model.py

@@ -2,6 +2,7 @@ import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
 
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-5):
@@ -33,7 +34,7 @@ class ResonanceLayerKaggle(nn.Module):
         phase = torch.tanh(self.to_phase(h)) * math.pi
 
         # Distance - causal ke liye
-        pos = torch.arange(T, device=x.device).float()
+        pos = torch.arange(T, device=x.device, dtype=x.dtype)
         chunk = 512  # safe for Kaggle T4
         out = torch.zeros(B, T, C, device=x.device, dtype=x.dtype)
         
@@ -49,7 +50,7 @@ class ResonanceLayerKaggle(nn.Module):
             angle = (phase_i - phase_j) + freq_i * dist * 0.05
             score = torch.cos(angle) / math.sqrt(self.head_dim)  # Normalize to prevent explosion
             
-            causal_mask = (torch.arange(i, end, device=x.device)[:,None] >= pos[None,:]).float().view(1,1,end-i,T)
+            causal_mask = (torch.arange(i, end, device=x.device)[:,None] >= pos[None,:]).to(x.dtype).view(1,1,end-i,T)
             score = score * causal_mask
             
             amp = A.norm(dim=-1).permute(0,2,1) # B,H,T
@@ -86,7 +87,7 @@ class ViuResonance100M(nn.Module):
     def forward(self, idx, targets=None):
         x = self.emb(idx)
         for layer in self.layers:
-            x = x + self.resid_dropout(layer(x))
+            x = x + self.resid_dropout(checkpoint(layer, x, use_reentrant=False))
         logits = self.head(self.norm(x))
         loss = None
         if targets is not None: