I want to replace average pooling in one architecture with DC part of DCT transformation, but when I replace this with each other I got nan values for loss.
This is my DCT transform over all channels:
def dct3(self, x):
    X1 = dct.dct_3d(x[0])
    X2 = dct.dct_3d(x[1])
    a = torch.stack([X1,X2], 3)
    for i in range((x.shape[0])-2):
        X = dct.dct_3d(x[i+2])
        X = X.unsqueeze(3)
        a = torch.cat([a,X], 3)
    a = a.reshape(x.shape[0],x.shape[1],7,7)
    a = a[:,:,1,1]
    a = a.reshape(x.shape[0],x.shape[1],1,1)
    return a   

And this is my forward path:
def forward(self, x):
    x = self.first_conv(x)
    x = self.second(x)
    x = self.features(x)
    x = self.conv_last(x)
    #x = self.globalpool(x)
    #x = self.dct3(x)
    if self.model_size == '2.0x':
        x = self.dropout(x)
    x = x.contiguous().view(-1, self.stage_out_channels[-1])
    x = self.classifier(x)
    return x