gpu support continue

[lazarusA]

this follows from #256

[gemini-code-assist]

Code Review

This pull request introduces GPU support and refactors the data pipeline to separate predictors and forcings, updating the training, evaluation, and checkpointing logic accordingly. Feedback identifies several critical issues, including a typo in the loss computation that will cause an UndefVarError, a statistically incorrect R² formula, and flawed NaN masking logic that could lead to incorrect loss values in multi-target models. Additionally, improvements are suggested for configuration type safety, the removal of leftover comments, and fixing a broken warning check in the data preparation logic.

[gemini-code-assist]

There is a typo on line 109: ŷ (y + U+0302) is used instead of the argument ŷ (U+0177), which will cause an UndefVarError. Additionally, y_nan should be target-specific (e.g., y_nan[target]) to ensure that NaNs in one target do not invalidate observations for other targets during loss calculation.

                y_t = y[target]
                ŷ_t = ŷ[target]
                _apply_loss(ŷ_t, y_t, y_nan[target], loss_spec)

[gemini-code-assist]

The $R^2$ formula is statistically incorrect. The denominator (Total Sum of Squares) should be calculated using the mean of the observed values y, not the predicted values ŷ.

    return 1 - sum((y[y_nan] .- ŷ[y_nan]).^2) / sum((y[y_nan] .- mean(y[y_nan])).^2)

[gemini-code-assist]

The current logic computes a global NaN mask by OR-ing all targets. This is problematic for multi-target models: if target A is NaN but target B is valid at index i, the mask will be true, and the loss function will attempt to compute a loss on the NaN value for target A, resulting in a NaN total loss. The mask should be kept per-target (e.g., as a NamedTuple of masks).

        is_no_nan = map(vec -> .!isnan.(vec), y)
        !any(map(any, is_no_nan)) && continue

[gemini-code-assist]

Similar to the training loop, the NaN masks for evaluation should be per-target to avoid cross-contamination of NaN values between different target variables.

    is_no_nan_t = map(vec -> .!isnan.(vec), y_train)
    is_no_nan_v = map(vec -> .!isnan.(vec), y_val)

[gemini-code-assist]

The fields gdev and cdev should have explicit type annotations (e.g., Lux.AbstractLuxDevice) to improve code clarity and potentially help with compiler optimizations. Additionally, ensure that gpu_device() is the intended default for all instances of DataConfig, as it may trigger device initialization.

    "Select a gpu_device or default to cpu if none available"
    gdev::Lux.AbstractLuxDevice = gpu_device()

    "Set the `cpu_device`, useful for sending back to the cpu model parameters"
    cdev::Lux.AbstractLuxDevice = cpu_device()

[gemini-code-assist]

Remove leftover comments and incomplete words.

[gemini-code-assist]

The warning check uses predictors_forcing, which is an empty Symbol[] initialized at line 89 and never updated. The logic should check if both predictors and forcings are empty instead.

    if isempty(predictors) && isempty(forcings)
        @warn "Note that you don't have predictors or forcing variables."
    end