def METHOD_NAME(self, **kwargs):
    if self._hvd:
        from tensorflow.python.distribute import hvd_strategy
        with hvd_strategy.METHOD_NAME() as context:
            yield context
    elif self._hb:
        with self._hb.METHOD_NAME() as context:
            yield context