[deepspeed] check whether model is NLP one instead of counting on input type
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index 079696c244fef..d34449fa57db9 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -2562,8 +2562,8 @@ def _prepare_input(self, data: Union[torch.Tensor, Any]) -> Union[torch.Tensor,
             return type(data)(self._prepare_input(v) for v in data)
         elif isinstance(data, torch.Tensor):
             kwargs = {"device": self.args.device}
-            if self.deepspeed and data.dtype != torch.int64:
-                # NLP models inputs are int64 and those get adjusted to the right dtype of the
+            if self.deepspeed and (torch.is_floating_point(data) or torch.is_complex(data)):
+                # NLP models inputs are int/uint and those get adjusted to the right dtype of the
                 # embedding. Other models such as wav2vec2's inputs are already float and thus
                 # may need special handling to match the dtypes of the model
                 kwargs.update({"dtype": self.args.hf_deepspeed_config.dtype()})

# What does this PR do?
This PR intends to fix an issue when training of NLP model fails if input dtype isn't int64.
My dataset had dtype = int32. Everything was ok until I decided to add deepspeed. 
It turned out that trainer relies on dtype and does input data convertion into hf_deepspeed_config.dtype if it isn't int64.
I guess it has to check whether first layer isn't Embedding instead. 

I think this PR also needs tests but I need an advice on how we can cover this case.  

@stas00 could you be so kind and review this PR and give an advice on whether tests are necessary and their implementation ?