Use Union in Torchscript

[oke-aditya]

🚀 The feature

Use Union in Torchscript.

pytorch/pytorch#64234

Motivation, pitch

Now that Union typing support is merged. We can rewrite the workarounds which don't support Union.
E.g we can now proceed with #4323

There maybe more places so we can track them in this Issue too.

Alternatives

No response

Additional context

No response

[NicolasHug]

That's great news, on top of #4323, a top priority regarding this would be to fix the annotations that are currently incorrect because we were lacking Union support in the past

[oke-aditya]

I think it is supported but the docs haven't updated the change?

https://pytorch.org/docs/master/jit_language_reference.html#supported-type

[frgfm]

I think it is supported but the docs haven't updated the change?

https://pytorch.org/docs/master/jit_language_reference.html#supported-type

@oke-aditya pytorch/pytorch#68435 should fix that 👍

[oke-aditya]

So nice of you @frgfm Francois. Merci beaucoup!!