I want to use autoencoder’s latent vector as feature extractor. For that, I want to put same input/output image for the autoencoder model.

I want to make “input_size” in both encoder & decoder as a same image.
Current model’s result is not a fixed. it just assigned the size of the output space.enter code here

import torch
import torch.nn as nn

# Define the autoencoder architecture
class Autoencoder(nn.Module):
    def __init__(self, input_size, latent_dim):
        super(Autoencoder, self).__init__()

        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(input_size, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, latent_dim)
        )

        # Decoder
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 32),
            nn.ReLU(),
            nn.Linear(32, 64),
            nn.ReLU(),
            nn.Linear(64, input_size)
        )

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