class ResidualAttentionTranspose(nn.Module):
    def __init__(self, dim):
        super(ResidualAttentionTranspose, self).__init__()
        self.linear1 = ResidualLinear(dim, dim)
        self.attenton = nn.MultiheadAttention(dim, 8, batch_first=True)
                

    def forward(self, x):
        shape = x.shape
        x = x.flatten(2).transpose(1, 2)
        x = self.linear1(x)
        attn_output, attn_weights = self.attenton(x, x, x)

        return (x + attn_output).view(shape)
    
class ResidualAttentionFlatten(nn.Module):
    def __init__(self, dim):
        super(ResidualAttentionFlatten, self).__init__()
        self.linear1 = ResidualLinear(dim, dim)
        self.attenton = nn.MultiheadAttention(dim, 8, batch_first=True)
                

    def forward(self, x):
        shape = x.shape
        x = x.flatten(2)
        x = self.linear1(x)
        attn_output, attn_weights = self.attenton(x, x, x)

        return (x + attn_output).view(shape)

class PrintShape(nn.Module):
    def __init__(self):
        super(PrintShape, self).__init__()

    def forward(self, x):
        #print(x.shape)
        return x

class PositionEmbedding(nn.Module):
    def __init__(self, dim):
        super(PositionEmbedding, self).__init__()
        self.register_buffer("position_ids", torch.arange(16384).expand((1, -1)), persistent=False)
        self.position_embedding = nn.Embedding(16384, dim)
        
    def forward(self, x):
        shape = x.shape
        
        x = x.flatten(2).transpose(1, 2)
        embeddings = self.position_embedding(self.position_ids)
        x =  x + embeddings

        return x.transpose(2, 1).view(shape)
		
encoder = nn.Sequential(
	nn.Conv2d(in_channels=3, out_channels=8, kernel_size=4, stride=2, padding=1),
    PrintShape(),
	#ResidualAttentionFlatten(64),
	nn.ReLU(),
	nn.Conv2d( in_channels=8, out_channels=8, kernel_size=3, stride=2, padding=1),
    PrintShape(),
	#ResidualAttentionFlatten(32*32),
	nn.ReLU(),    
	nn.Conv2d( in_channels=8, out_channels=8, kernel_size=3, stride=2, padding=1),
    PrintShape(),
	#ResidualAttentionFlatten(1024),
	nn.Conv2d( in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1),
	#ResidualAttentionFlatten(256),
    PrintShape(),
	#ResidualAttentionFlatten(64*64),
	nn.Conv2d( in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1),
    PrintShape(),
	ResidualAttentionFlatten(64*64),
    #PrintShape(),
	).to('cuda')


#linear = ResidualAttention(768)
decoder = nn.Sequential(
	#PrintShape(),
    nn.ConvTranspose2d(in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1),
	#ResidualAttentionFlatten(8),
	#nn.ConvTranspose2d(in_channels=384, out_channels=256, kernel_size=4, stride=2, padding=1),
	#ResidualAttentionFlatten(1024),
    PrintShape(),
	#ResidualAttentionTranspose(256),
	nn.ConvTranspose2d(in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1),
    PrintShape(),
	nn.ReLU(),
	#ResidualAttentionTranspose(128),
	nn.ConvTranspose2d(in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1),
	#ResidualAttentionTranspose(32),
    PrintShape(),
	nn.ConvTranspose2d(in_channels=8, out_channels=8, kernel_size=4, stride=2, padding=1 ),
	#ResidualAttentionTranspose(16),
    PrintShape(),
	nn.ConvTranspose2d(in_channels=8, out_channels=8, kernel_size=4, stride=2, padding=1 ),
    #nn.ReLU(),
    PrintShape(),
	nn.ConvTranspose2d(in_channels=8, out_channels=3, kernel_size=4, stride=2, padding=1 ),
	#nn.ConvTranspose2d(in_channels=16, out_channels=3, kernel_size=4, stride=2, padding=1 ),

	).to('cuda')


why it cant reconsutrct colors even after 200k steps