Hello, after experimenting with multiple off-the-shelf and written from scratch networks I am starting to feel there is something wrong with my network without being able to understand what:
My network
class MySubPixelCNN(nn.Module):
    def __init__(self, upscale_factor,num_features):
        super(MySubPixelCNN, self).__init__()

        self.relu = nn.ReLU()
        self.conv1 = nn.Conv2d(num_features, 64, kernel_size=5, stride=1, padding=2)
        self.bb1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
        self.bb2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1)
        self.bb3 = nn.BatchNorm2d(32)
        self.conv4 = nn.Conv2d(32, num_features*upscale_factor ** 2, kernel_size=3, stride=1, padding=1)
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bb1(self.relu(x))
        x = self.conv2(x)
        x = self.bb2(self.relu(x))
        x = self.conv3(x)
        x = self.bb3(self.relu(x))
        x = self.conv4(x)
        x = self.pixel_shuffle(x)
        return x

My input images are textures (I.e not the ordinary Images), with a simply preprocessing to have them in [0,1].
Yet, sometimes (which makes the network immediately diverge)
model(input).max() > 1e30
As mentioned, I always add batch norm, and as the network is not that deep, I simply cannot understand what could’ve gone wrong.