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[0]):
            
            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