Change the device to execute on for the scope of the block.

def use_device(cls, device):
        """Change the device to execute on for the scope of the block."""
        if device == cls.ops.device:
            yield
        else:
            curr_Ops, curr_ops = (cls.Ops, cls.ops)
            cls.Ops = get_ops(device)
            cls.ops = cls.Ops()
            yield
            cls.Ops = curr_Ops
            cls.ops = curr_ops