Fix causal mask dtype for float16

app.py

@@ -75,6 +75,7 @@ class GcodeDecoder(nn.Module):
     def forward(self, latent: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
         batch_size, seq_len = input_ids.shape
         device = input_ids.device
+        dtype = latent.dtype
         
         latent_flat = latent.view(batch_size, -1)
         memory = self.latent_proj(latent_flat)
@@ -83,7 +84,8 @@ class GcodeDecoder(nn.Module):
         positions = torch.arange(seq_len, device=device)
         x = self.token_embed(input_ids) + self.pos_embed(positions)
         
-        causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device)
+        # Causal mask must match dtype for attention
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device, dtype=dtype)
         
         for layer in self.layers:
             x = layer(x, memory, tgt_mask=causal_mask)