Make gradient checkpointing configurable in the predict

[javak87]

Description

In order to save memory and improve computational efficiency, and since gradient checkpointing is used in several parts of the code, having the ability to enable or disable gradient checkpointing in specific sections can provide flexibility to maximize performance while also conserving memory.

Issue Number

Refs #1141

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

Comparing the baseline and current PR performance

When pred_gradient_checkpoint_mode is set to false, the performance and GPU memory peak are as follows:

default_config.yml:

../WeatherGenerator-private/hpc/launch-slurm.py --time 60

mixed.yml:

../WeatherGenerator-private/hpc/launch-slurm.py --time 60 --config ./config/mixed.yml

For those GPUs with memory more than 18 GiB, it is recommended to set pred_gradient_checkpoint_mode to False