criterion.pos_weight = torch.tensor([loss_fn_pos_weight])
The above line ‘criterion.pos_weight’ is not defined on GPU, it should be written like this:
criterion.pos_weight = torch.tensor([loss_fn_pos_weight], device=device)