Thank you for the fast reply.
I have managed to get my autoencoder “right”  , I have a bottleneck layer of 2x2x2 which gives me 8 clusters like this:
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 2, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.MaxPool2d(6, stride=1)
        )

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(2, 64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 1, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

and my training is simple, the loss is how much the prediction differs from the original image:

    distance = nn.MSELoss()

    loss = distance(states,output)

but now the problem is that the network is not learning (the loss is constant at 10k which is bad)
even with a high learning rate:
optimizer_encoder = optim.SGD(model2.parameters(), lr=0.005)
I have no idea why