I’ve put together a small example hoping it could help find the issue, but everything works as expected:
image785×1086 88.9 KB
I created a simplified architecture, still no go:
class FrontEnd(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_parameter('alpha', torch.nn.Parameter(torch.tensor(0.)))
        
    def forward(self, x):
        x = x ** torch.sigmoid(self.alpha) # <----- line causing issues
        return x

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.frontend = FrontEnd()
        self.lin = nn.Linear(508800, len(classes))
    
    def forward(self, x):
        bs, im_num, ch, y_dim, x_dim = x.shape
        x = self.frontend(x)
        return self.lin(x.view(16, -1))

image1938×1184 221 KB
I wonder what I could try to fix this or troubleshoot this further? How can I pinpoint why PowBackward1 returned nan values?