Fix DDP parameter and gradient layouts

models.py

@@ -100,8 +100,11 @@ class DeepHealth(nn.Module):
         self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
         if target_mode == "next_token":
             self.risk_head.weight = self.token_embedding.weight
-        self.query_token = nn.Parameter(torch.zeros(n_embd))
-        nn.init.normal_(self.query_token, mean=0.0, std=0.02)
+        if target_mode == "all_future":
+            self.query_token = nn.Parameter(torch.zeros(n_embd))
+            nn.init.normal_(self.query_token, mean=0.0, std=0.02)
+        else:
+            self.register_parameter("query_token", None)
 
     def _make_history_attn_mask(
         self,
@@ -251,6 +254,11 @@ class DeepHealth(nn.Module):
         h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 
         if mode == "all_future":
+            if self.query_token is None:
+                raise RuntimeError(
+                    "all_future forward requires a model initialized with "
+                    "target_mode='all_future'"
+                )
             batch_size = event_seq.size(0)
             query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
             h_disease = torch.cat([h_disease, query], dim=1)