The problem is that pos_weight is not actually differentiable because the argmax op is not differentiable.
But due to a  bug on our side, this is detected too late and here is a minimal repro:
import torch
from torch import nn

t = torch.rand(10, requires_grad=True)
bad = torch.argmax(t)

res = bad + 2

res.sum().backward()

You can fix your code by adding a .detach() to the output of torch.where before returning it in pos_weight.