Update train_curriculum.py

train_curriculum.py

@@ -47,11 +47,12 @@ except ImportError as _e:
     BaseCallback    = object
 
 try:
-    from gru_weather_policy import create_gru_weather_policy_kwargs
+    from gru_weather_policy import create_gru_weather_policy_kwargs, get_equivariant_policy_class
     _GRU_AVAILABLE = True
 except ImportError:
     _GRU_AVAILABLE = False
     create_gru_weather_policy_kwargs = None
+    get_equivariant_policy_class = None
 
 try:
     from physics_dynamics import TemporalDynamicsModel, DynaRolloutBuffer, ZoneStateTensor
@@ -437,7 +438,8 @@ class DynaCallback(BaseCallback):
                 total_loss = 0.0
                 n_batches  = 0
 
-                for i in range(0, len(pairs) - batch_size, batch_size):
+                # FIX: iterate over all pairs, including the final partial batch
+                for i in range(0, len(pairs), batch_size):
                     batch = pairs[i : i + batch_size]
                     curr_list = [p[0] for p in batch]
                     next_list = [p[1] for p in batch]
@@ -510,7 +512,11 @@ def _build_dyna_callback(
         return None
 
     try:
-        dynamics_model = TemporalDynamicsModel.load(str(model_path))
+        # FIX: load onto the same device as training to avoid CPU/CUDA mismatch
+        import torch as _torch
+        dynamics_model = TemporalDynamicsModel.load(
+            str(model_path), device=_torch.device(device)
+        )
         dynamics_model.eval()
         dyna_buffer = DynaRolloutBuffer(
             dynamics=dynamics_model,
@@ -636,9 +642,10 @@ def train_phase(
 
     env = Monitor(make_weather_env(config))
 
+    # FIX: use the zone-equivariant policy class instead of the generic string
     if _GRU_AVAILABLE:
         policy_kwargs = create_gru_weather_policy_kwargs(hidden_size=hidden_size)
-        policy = "MultiInputPolicy"
+        policy = get_equivariant_policy_class()
         logger.info("Using GRU policy (hidden_size=%d)", hidden_size)
     else:
         policy_kwargs = dict(net_arch=dict(pi=[128, 64], vf=[128, 64]))