[typing] TypeVar for broadcast rather than object

[kandluis]

The type for broadcast should probably be defined using TypeVar, rather than object since we only want to enforce that the input type is the same as the output type.

Something like:

T = TypeVar("T", bound=object)
def broadcast(self, obj: T, src: int = 0) -> T:
   ...

This would avoid requiring cast at the call-site when broadcasting elements.

https://github.com/PyTorchLightning/pytorch-lightning/blob/f79a13e4956554eb68b103c12054bdb7cd4dbe1a/pytorch_lightning/plugins/training_type/training_type_plugin.py#L90

[tchaton]

Dear @kandluis,

Great catch !
Would you mind making a PR to update the typing ?

Best,
T.C