But does that mean that if the device is not initialized yet then the current device will never be set on the worker thread? That sounds bad no?

Current device will always be initialized at this instance, because prior to this call PyTorch would set currents device for memory allocations and kernel calls in forward pass.

It's possible to run a custom autograd function that would run afoul of this, but we already operate under assumption that no new contexts are going to be created in autograd:

I'm still confused.
The following sounds valid but the second backward call will use worker threads that are not properly on the right device.:

import torch

a = torch.rand(10, requires_grad=True)
a.sum().backward() # This call creates the worker threads, no cuda context here

b = torch.rand(10, device="cuda", requires_grad=True)
b.sum().backward() # This one will now run with no current device set!!

Ok then would setting device for every thread somewhere here

You mean when you run it locally?

Both locally and CI

Well I do expect the code to run but you will need a bit more meat there to check that you're on the right device.
I don't have a multi-gpu machine right now but my guess is that the following will print all 0s instead of a 1 for the second BW:

import torch

class MyFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        print("FW", torch.cuda.current_device())
        return x.clone()

    @staticmethod
    def backward(ctx, gO):
        print("BW", torch.cuda.current_device())
        return gO

a = torch.rand(10, requires_grad=True)
MyFn.apply(a).sum().backward() # This call creates the worker threads, no cuda context here

b = torch.rand(10, device="cuda:1", requires_grad=True)
MyFn.apply(b).sum().backward() # This one will now run with no current device set!!

Curious to see what the result is if you have a multi-gpu machine at hand.

Works fine locally:

FW 0
BW 0
FW 1
BW 1

# And if I print them
# a is tensor([0.2930, 0.1664, 0.7401, 0.5832, 0.2043, 0.9405, 0.6830, 0.8309, 0.1875, 0.8577], requires_grad=True)
# a.grad is tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

# b is tensor([0.1209, 0.5240, 0.6959, 0.3283, 0.1513, 0.7672, 0.0403, 0.4019, 0.3628, 0.5973], device='cuda:1', requires_grad=True)
# b.grad is tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], device='cuda:1')

It could work by virtue of device guards, and the threads default device outside guards could be 0.