Hi Bran!



 stark:

I gave MSELoss a try using the actual class counts as below:
criterion = torch.nn.BCELoss().cuda()

loss = tensor(nan, device='cuda:0', grad_fn=<BinaryCrossEntropyBackward>)



The immediate problem is that you are actually using BCELoss
(binary cross-entropy), rather than MSELoss (mean-squared-error)
as your loss function.
BCELoss is not conceptually appropriate for your use class-count
use case.  Also, technically, it requires its target values to be
probabilities – numbers between 0.0 and 1.0 – for its computations
to make any sense, and your class-count targets are often greater
than 1.
Best.
K. Frank