How to create simple images?

vision
1

I was trying to build a very simple model to create simple circle images with 5 different colors and random radius, but I did not get a good result. I have created 500 images for each color for the training set.
I am just wondering that the logic of the model is correct.

The input shape is : one_hot([1,0,0,0,0]) and the output shape is (35,35,3).
Here is my simple model:

cls_model(
  (model): Sequential(
    (0): Linear(in_features=5, out_features=1024, bias=True)
    (1): ReLU()
    (2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
    (3): Linear(in_features=1024, out_features=961, bias=True)
    (4): ReLU()
  )
  (tconv1): ConvTranspose2d(1, 1, kernel_size=(3, 3), stride=(1, 1))
  (tconv2): ConvTranspose2d(1, 3, kernel_size=(3, 3), stride=(1, 1))
)

criterion = nn.MSELoss()
optim = torch.optim.Adam(net.parameters(), lr=1e-3)