Add t0 loss target

This commit is contained in:
Jaret Burkett committed 2026-03-28 13:35:21 -06:00
1 parent 8302b21f8f
commit 6a1fc54779
2 files changed
+18

No files matched your search

+3
View File
@@ -497,6 +497,9 @@ class TrainConfig:
self.correct_pred_norm_multiplier = kwargs.get('correct_pred_norm_multiplier', 1.0)
self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, pixelspace, mean_flow
# do the loss on a timestep to 0 prediction
self.t0_loss_target = kwargs.get('t0_loss_target', False)
# scale the prediction by this. Increase for more detail, decrease for less
self.pred_scaler = kwargs.get('pred_scaler', 1.0)