Update Losses to Gen 2

[djgagne]

Losses need to be updated to accept Gen 2 state dict, handle computed_diagnostics variables. MSE definitely should follow this framework. Other losses could potentially be added in 2026.2.1.

[djgagne]

Currently the call to the loss in trainer_gen2 takes y_pred as input,which is the raw output from the model.

miles-credit/credit/trainers/trainer_gen2.py

Lines 175 to 186 in 595cf70

	full_data_dict = apply_postblocks(self.step_postblocks, full_data_dict)
	if t in self.backprop_on_timestep:
	if self.flag_clamp:
	full_data_dict["y"] = torch.clamp(
	full_data_dict["y"].float(), min=self.clamp_min, max=self.clamp_max
	)
	with torch.autocast(device_type="cuda", enabled=self.amp):
	loss = criterion(
	full_data_dict["y"].float().to(full_data_dict["y_pred"].dtype),
	full_data_dict["y_pred"],
	).mean()

, The loss should be calculated against full_data_dict["y_processed"], and we need to make sure that all of the post blocks are updating the state in full_data_dict["y_processed"], which may not be happening with geopotential and mslp at the moment.

There is also an open question about how to process the target data for the model. Does it need to go through both the preblocks and postblocks? Does it need to be renormalized before the loss or have some kind of normalization or variance/covariance as part of the loss?