[tiny] Fix gradient checkpointing for Oumi trainer (#1778)

Author: wizeng23

src/oumi/core/trainers/oumi_trainer.py

@@ -30,7 +30,6 @@
 import torch.utils.tensorboard as tensorboard
 
 import mlflow  # isort: skip
-import transformers
 
 import wandb  # isort: skip
 from torch.distributed.checkpoint.state_dict import (
@@ -161,11 +160,11 @@ def __init__(
         # Prepare model for training
         # ----------------------------------
         if args.enable_gradient_checkpointing:
-            if not isinstance(model, transformers.PreTrainedModel):
+            if not hasattr(model, "gradient_checkpointing_enable"):
                 raise ValueError(
-                    "Gradient checkpointing is only supported for transformers models."
+                    "Gradient checkpointing is only supported for Hugging Face models."
                 )
-            model.gradient_checkpointing_enable(args.gradient_checkpointing_kwargs)
+            model.gradient_checkpointing_enable(args.gradient_checkpointing_kwargs)  # type: ignore
         model = cast(torch.nn.Module, model)
         model.to(self.device)
         if is_distributed():