Note that for all tensor creation functions without explicit values (empty,zero,full,randn etc.) host-device transfer (.cuda()) is much slower than device=X argument to initial function. Also, I think you can change manually_set_device from ‘none’ to None and avoid if-else checks. I.e.
cur_device = None
#@contextmanager doing DI
...
def make_tensor(shape):
  return torch.zeros(shape, device=cur_device)