Bug in my custom loss function still there after upgrading to torch 1.8.1

username: Sevenfold
cooked: <p>I have written a custom loss function and keep getting the error :</p>
<blockquote>
<p>UnboundLocalError: local variable ‘beta1’ referenced before assignment</p>
</blockquote>
<p>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.</p>
<p>Below is the function which is throwing the error.</p>
<pre><code class="lang-auto">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
</code></pre>

username: KFrank
cooked: <p>Hi Sevenfold!</p>
<aside class="quote no-group quote-modified" data-username="Sevenfold" data-post="1" data-topic="137235" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sevenfold/48/29724_2.png" class="avatar"> Sevenfold:</div>
<blockquote>
<p>I have written a custom loss function and keep getting the error :</p>
<blockquote>
<p>UnboundLocalError: local variable ‘beta1’ referenced before assignment</p>
</blockquote>
</blockquote>
</aside>
<p>This does seem to be a known bug and the <a href="https://github.com/pytorch/pytorch/pull/52944" rel="noopener nofollow ugc">relevant pull request</a><br>
does seem to suggest that the bug was fixed in 1.8.1.</p>
<p>Nevertheless, perhaps the fix didn’t make it into 1.8.1, or perhaps<br>
you’ve uncovered a different bug that looks similar.</p>
<p>In any event, could you post a small, self-contained, fully-runnable<br>
script that reproduces the issue?  You might also try upgrading to<br>
the latest stable version, 1.10.0, just to see (and maybe, just for fun,<br>
to the latest nightly).</p>
<p>Best.</p>
<p>K. Frank</p>