Special handling of Torch DDP in callback (#21081)

* Special handling of Torch DDP in callback

* Use inheritance tree for DDP check

Modified DDP check to use isinstance rather than type().__name__ for
robustness. Fixed additional whitepace

* Fixing comment.

* inlining DDP import where its needed.

Author: timsweeneyfanelli

keras/src/callbacks/callback.py

@@ -76,6 +76,19 @@ def set_model(self, model):
 
     @property
     def model(self):
+        if backend.backend() == "torch":
+            from torch.nn.parallel import DistributedDataParallel
+
+            if isinstance(self._model, DistributedDataParallel):
+                # Keras Callbacks expect to work with Keras models. e.g
+                # ModelCheckpoint and EarlyStopping both attempt to call
+                # keras-specific APIs on the value returned from this
+                # property. If this callback was created against a DDP
+                # wrapper instead of the underlying keras.Model, it is
+                # likely to fail. Return self._model.module for DDP
+                # instances instead.
+                return self._model.module
+
         if backend.backend() == "jax" and hasattr(
             self._model, "jax_state_sync"
         ):