fix teacher forcing bug

OuyangWenyu · Nov 7, 2024 · f5f97e1 · f5f97e1
1 parent d879f3f
commit f5f97e1
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/torchhydro/models/seq2seq.py b/torchhydro/models/seq2seq.py
@@ -187,7 +187,13 @@ def forward(self, *src):
             use_teacher_forcing = (
                 random_vals < self.teacher_forcing_ratio
             ) * valid_mask
-            current_input = trg * use_teacher_forcing + output * (~use_teacher_forcing)
+            current_input = torch.where(
+                torch.isnan(trg),  # if trg is nan
+                output,  # then use output
+                trg * use_teacher_forcing
+                + output
+                * (~use_teacher_forcing),  # else calculate with teacher forcing
+            )
 
         outputs = torch.stack(outputs, dim=1)
         if self.prec_window > 0: