Address reviewer feedback: rollout bug, dead amp field, _step_ensemble indirection

- Fix inference _rollout: torch.cat([history[:, 1:], next_frame.unsqueeze(1)])
  so history window slides correctly for any history_frames value, not just 2
- Remove unimplemented amp config field from TrainingConfig and default.yaml
- Inline model call in _loss AR loop instead of routing through _step_ensemble
  with num_samples=1 (each member needs its own history, so the single-call
  collapse doesn't apply; direct call is cleaner)

Signed-off-by: Kashif Rasul <kashif.rasul@gmail.com>

Author: kashif

examples/weather/fgn/config/training/default.yaml

@@ -12,7 +12,6 @@ checkpoint_freq: 50
 validation_freq: 25
 resume_checkpoint: latest
 clip_grad_norm: -1.0
-amp: false
 
 # Autoregressive rollout depth per training step.
 # Must match dataset.future_frames. 1 = single-step training (paper Stages

examples/weather/fgn/inference.py

@@ -106,8 +106,8 @@ def _rollout(
                 next_frame = next_frame.clone()
                 for ci in output_only_channels:
                     next_frame[:, ci].zero_()
-            rollout_history = torch.stack(
-                [rollout_history[:, 1], next_frame],
+            rollout_history = torch.cat(
+                [rollout_history[:, 1:], next_frame.unsqueeze(1)],
                 dim=1,
             )
         trajectories.append(torch.stack(states, dim=1))

examples/weather/fgn/utils/config.py

@@ -58,7 +58,6 @@ class TrainingConfig:
     validation_freq: int = Field(default=25, ge=1)
     resume_checkpoint: int | Literal["latest"] | None = "latest"
     clip_grad_norm: float = -1.0
-    amp: bool = False
     ar_steps: int = Field(default=1, ge=1, le=8)
     # Data + domain parallelism knobs. Mirrors StormCast's convention.
     # - domain_parallel_size=1 & force_sharding=False → pure single-process

examples/weather/fgn/utils/trainer.py

@@ -320,13 +320,20 @@ def _loss(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
             members = []
             for n in range(num_samples):
                 hist_n = per_member_hist[:, n]
-                pred_n = self._step_ensemble(
-                    history=hist_n,
-                    background=background,
-                    invariants=invariants,
-                    num_samples=1,
-                )[:, 0]
-                members.append(pred_n)
+                latent = torch.randn(
+                    hist_n.shape[0],
+                    int(self.cfg.model.latent_dim),
+                    device=self.device,
+                    dtype=torch.float32,
+                )
+                members.append(
+                    self.model(
+                        history=hist_n,
+                        latent=latent,
+                        background=background,
+                        invariants=invariants,
+                    )
+                )
             preds = torch.stack(members, dim=1)  # (B, N, C, H, W)
 
             step_loss = fair_crps(preds, target[:, k], weights=self.loss_weights)