diff --git a/src/axolotl/utils/schemas/config.py b/src/axolotl/utils/schemas/config.py
index 1ec9e296d..a0d04ccc2 100644
--- a/src/axolotl/utils/schemas/config.py
+++ b/src/axolotl/utils/schemas/config.py
@@ -1290,3 +1290,5 @@ class AxolotlConfigWCapabilities(AxolotlInputConfig):
                 LOG.warning(
                     f"torch=={torch_version} may not be supported in future versions. Please consider upgrading to torch>=2.5.1."
                 )
+
+        return self