import torch import torch.nn as nn class Autoencoder(nn.Module): def __init__(self, input_dim): super(Autoencoder, self).__init__() self.encoder = nn.Sequential( nn.Linear(input_dim, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU() ) self.decoder = nn.Sequential( nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, input_dim) ) def forward(self, x): encoded = self.encoder(x) decoded = self.decoder(encoded) return decoded