When are buffers moved to gpu?
I have an issue with a weighted mse function that I instantiate in the setup, with a buffer as parameter. Something like this:
@torch.jit.script
def weighted_mse_func(weights, y, y_hat):
    # weighted regression loss
    reg_loss = torch.dot(weights,torch.mean(F.mse_loss(y_hat, y, reduction='none'), dim=0))
    return reg_loss

def weighted_mse(weights):
    def func(y, y_hat):
        return weighted_mse_func(weights, y, y_hat)
    return func


class model(pl.LightningModule):
    def __init__(self, weights):
        weights = torch.tensor(weights.copy(), dtype=self.dtype, device=self.device)
        self.register_buffer("weights", weights)
    
    def setup(self, stage):
        super().setup(stage)
        self.loss = weighted_mse(self.weights)
When initializing training on the GPU I get an error because self.weights is on CPU and not in GPU, if after the error I check the device of the buffer it's on GPU. So if I re-run the trainer, it works fine, also works fine if I call model.cuda() before training. What is going on? Why is the buffer not in GPU on the setup where it fails, but it is afterward? Is something asynchronous going on here?
It seems that pytorch doesn't move buffers parameters in-place (like it is done for parameters), this results in references to buffers being useless if they are moved from one device to another. This issue is discussed in pytorch/pytorch#43815.