Fabric leaks the default device on exception

[carmocca]

Bug description

See reproduction below

What version are you seeing the problem on?

master

How to reproduce the bug

import pytest
import torch

import lightning as L
from lightning.fabric.plugins import HalfPrecision


class MyPlugin(HalfPrecision):
    def init_context(self):
        # arbitrary error
        raise NotImplementedError


fabric = L.Fabric(devices=1, accelerator="cuda", plugins=MyPlugin("16-true"))

with pytest.raises(NotImplementedError):
    with fabric.init_module():
        pass

default_device = torch.tensor(0).device
assert default_device.type == "cpu", default_device

Error messages and logs

The assertion fails

Environment

Current master

More info

This blocks #18704

[awaelchli]

Could this be resolved by resetting the device in a finally-block?

[carmocca]

I thought torch.device implemented __exit__. If not, then we would need to write our own context manager that does it

[awaelchli]

The PyTorch context manager does it correctly, but since we are using exit stack this functionality gets lost. This reproduces:

from contextlib import ExitStack
import pytest
import torch

stack = ExitStack()
ctx = stack.enter_context(torch.device("cuda:0"))
# ctx = torch.device("cuda:0")  # WORKS!

with pytest.raises(NotImplementedError):
    with ctx:
        raise NotImplementedError()

default_device = torch.tensor(0).device
assert default_device.type == "cpu", default_device

But if you remove the exit stack and just use torch.device() directly, it passes.

[carmocca]

Your example should be

from contextlib import ExitStack
import pytest
import torch

ctx = ExitStack()
ctx.enter_context(torch.device("cuda:0"))

with pytest.raises(NotImplementedError):
    with ctx:
        raise NotImplementedError()

default_device = torch.tensor(0).device
assert default_device.type == "cpu", default_device

which does work as expected.

I found the cause of this issue. I'll open a PR