Update code/model.py with chunked cross-entropy and H200 turbo speed optimizations

code/model.py

@@ -195,7 +195,7 @@ class ViuAI(nn.Module):
                 tgt = targets
 
             total_loss = torch.zeros(1, device=x.device, dtype=torch.float32)
-            total_active = 0
+            total_active = (tgt != -100).sum().clamp(min=1)
             seq_len = hidden.size(1)
 
             for i in range(0, seq_len, ce_chunk):
@@ -206,11 +206,9 @@ class ViuAI(nn.Module):
                     chunk_logits.view(-1, V), chunk_tgt.view(-1),
                     ignore_index=-100, reduction="sum"
                 )
-                chunk_active = (chunk_tgt != -100).sum().item()
                 total_loss = total_loss + chunk_loss
-                total_active += chunk_active
 
-            loss = total_loss / max(1, total_active)
+            loss = total_loss / total_active
 
             # z-loss: stabilizes logit magnitudes during pretraining; disable for SFT (z_loss_weight=0)
             if self.cfg.z_loss_weight > 0: