I have written a custom loss function and keep getting the error :
UnboundLocalError: local variable ‘beta1’ referenced before assignment
The error occurs when I use my loss function (see below) along with AdamW as my optimizer. According to the forums this was caused by a bug in PyTorch 1.8.0, so I upgraded to 1.8.1 but still the bug remains.
Below is the function which is throwing the error.
class SoftCel(nn.Module): def __init__(self, ignore_val, class_weights = None, device = 'cpu'): super(SoftCel, self).__init__() self.device = device self.ignore = ignore_val self.class_weights = class_weights.to(self.device) if self.class_weights is None: self.class_weights = torch.ones((5, 1, 1)).to(self.device) def get_class_weights_for_mean(self, labels): """If we use class weights we need to normalize by the weights as per pytorch forum: loss_weighted_manual = loss_weighted_manual.sum() / weights[target].sum() Thus this function gets us the denominator of the above EQ. """ weights =  classes = torch.unique(labels) for c in classes: if c == self.ignore: continue # number times a label is present num = len(labels[labels == c]) # get weight value for that label w = self.class_weights.squeeze()[c].item() temp = [w for _ in range(num)] weights.extend(temp) return weights def forward(self, preds, labels, one_hot, pixel_weights = None): # list for pixel loss values and pixel weight values L =  W =  for i in range(preds.shape): loss_val = - one_hot[i].long() * F.log_softmax(preds[i]) if pixel_weights is not None: loss_val = loss_val * pixel_weights[i] if self.class_weights is not None: # #loss_val = self.class_weights * loss_val loss_val = loss_val * self.class_weights loss_val = torch.sum(loss_val, 0) if self.ignore is not None: loss_val = loss_val[labels[i]!=self.ignore] L.extend(loss_val) W.extend(self.get_class_weights_for_mean(labels[i])) L = torch.tensor(L) W = torch.tensor(W) # get mean loss = L.sum() / W.sum() loss.requires_grad = True return loss