[trainer] conditional ctx managers into one wrapper
This PR simplifies the code by taking 3 conditionals that are repeated in 3 places in trainer.py and makes use of a wrapper method.

Let me know if you prefer I make this helper method private and add a leading `_`.

If this looks good, we could repeat this in the other places where we have similar situations. While this makes the code simpler it doesn't necessarily make it more easy to understand, since it makes some parts of the logic somewhat hidden.

@sgugger 

