hey, i was planning to make a image generator especially for dsa, so like let's say i have to make aarray of n size and so so i wanted to make a cgan for that
i'll provide you with the code
i am using kaggle for this
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np
import os
import random
from PIL import Image
import json
import matplotlib.colors as mcolors  # Correct import for colors

def draw_3d_array(array_values, 
                 title="Array Visualization",
                 index_color='black', 
                 value_color='white',
                 show_labels=True,
                 elevation=30, 
                 azimuth=30,
                 save_path=None):
    """
    Draws a 3D visualization of an array with enhanced realism.
    """
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')

    # Parameters
    cube_size = 0.8  # Slightly smaller cubes for spacing
    spacing = 1.0     # Reduced spacing

    if not array_values:
        # Handle empty array: Display a message
        ax.text(0, 0, 0, "Empty Array", color='red', ha='center', va='center', fontsize=14)
    else:
        for i, val in enumerate(array_values):
            # Define the starting corner of the cube
            r = [i * spacing, 0, 0]
            # Random slight variation in position to add realism
            r_variation = [random.uniform(-0.1, 0.1) for _ in range(3)]
            r = [r[j] + r_variation[j] for j in range(3)]
            points = np.array([
                [r[0], r[1], r[2]],
                [r[0] + cube_size, r[1], r[2]],
                [r[0] + cube_size, r[1] + cube_size, r[2]],
                [r[0], r[1] + cube_size, r[2]],
                [r[0], r[1], r[2] + cube_size],
                [r[0] + cube_size, r[1], r[2] + cube_size],
                [r[0] + cube_size, r[1] + cube_size, r[2] + cube_size],
                [r[0], r[1] + cube_size, r[2] + cube_size]
            ])

            # Define the vertices for each face of the cube
            vertices = [
                [points[0], points[1], points[2], points[3]],  # Bottom face
                [points[4], points[5], points[6], points[7]],  # Top face
                [points[0], points[1], points[5], points[4]],  # Front face
                [points[2], points[3], points[7], points[6]],  # Back face
                [points[1], points[2], points[6], points[5]],  # Right face
                [points[4], points[7], points[3], points[0]]   # Left face
            ]

            # Assign a random color from a predefined palette to add diversity
            color_palette = [
                'skyblue', 'lightgreen', 'salmon', 'plum', 'lightcoral',
                'lightseagreen', 'paleturquoise', 'lightpink', 'lightsalmon',
                'mediumslateblue'
            ]
            selected_color = random.choice(color_palette)
            # Convert color to RGBA if needed for alpha
            if isinstance(selected_color, str):
                selected_color_rgb = mcolors.to_rgb(selected_color)
            else:
                selected_color_rgb = selected_color
            face_color = selected_color_rgb + (0.9,)  # Adding alpha for slight transparency

            # Create the cube
            poly3d = Poly3DCollection(vertices, facecolors=face_color, linewidths=0.5, edgecolors='k', alpha=0.9)
            ax.add_collection3d(poly3d)

            # Add index and value labels with slight randomization in position
            if show_labels:
                # Index label below the cube
                ax.text(r[0] + cube_size / 2, r[1] - 0.3, r[2], f"{i}", color=index_color,
                        ha='center', va='top', fontsize=8, weight='bold')
                # Value label inside the cube
                ax.text(r[0] + cube_size / 2, r[1] + cube_size / 2, r[2] + cube_size / 2, str(val), 
                        color=value_color, ha='center', va='center', fontsize=10, weight='bold')

    # Set the axes limits based on array size
    if not array_values:
        ax.set_xlim(-2, 2)
    else:
        ax.set_xlim(-1, len(array_values) * spacing)
    ax.set_ylim(-1, 2)
    ax.set_zlim(-1, 2)

    # Remove axis ticks and panes
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    ax.xaxis.pane.fill = False
    ax.yaxis.pane.fill = False
    ax.zaxis.pane.fill = False

    # Set viewing angle
    ax.view_init(elev=elevation, azim=azimuth)

    # Set title
    plt.title(title, fontsize=16, pad=20)

    # Set background color
    ax.set_facecolor('white')

    # Adjust layout
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Image saved to {save_path}")
    else:
        plt.show()

    # Clear the figure to free memory
    plt.clf()

def draw_infinity_array(save_path="infinity_array.png"):
    """
    Draws a conceptual 3D visualization representing an infinite array.
    Enhanced with varied colors and more realistic shading.
    """
    fig = plt.figure(figsize=(12, 7))
    ax = fig.add_subplot(111, projection='3d')

    # Parameters
    cube_size = 0.8
    spacing = 1.0
    num_cubes = 50  # Number of cubes to simulate infinity
    fade_start = 30  # Starting point for fade-out effect

    # Color palette for realism
    color_palette = [
        'lightcoral', 'tomato', 'darkorange', 'mediumpurple',
        'steelblue', 'seagreen', 'gold', 'orchid', 'lightseagreen', 'salmon'
    ]

    for i in range(num_cubes):
        # Define the starting corner of the cube, arranged diagonally for a 3D effect
        r = [i * spacing * 1.2, i * spacing * 0.6, 0]
        points = np.array([
            [r[0], r[1], r[2]],
            [r[0] + cube_size, r[1], r[2]],
            [r[0] + cube_size, r[1] + cube_size, r[2]],
            [r[0], r[1] + cube_size, r[2]],
            [r[0], r[1], r[2] + cube_size],
            [r[0] + cube_size, r[1], r[2] + cube_size],
            [r[0] + cube_size, r[1] + cube_size, r[2] + cube_size],
            [r[0], r[1] + cube_size, r[2] + cube_size]
        ])

        # Define the vertices for each face of the cube
        vertices = [
            [points[0], points[1], points[2], points[3]],  # Bottom face
            [points[4], points[5], points[6], points[7]],  # Top face
            [points[0], points[1], points[5], points[4]],  # Front face
            [points[2], points[3], points[7], points[6]],  # Back face
            [points[1], points[2], points[6], points[5]],  # Right face
            [points[4], points[7], points[3], points[0]]   # Left face
        ]

        # Apply a fade-out effect for cubes beyond fade_start
        if i >= fade_start:
            alpha = max(0.1, 1 - (i - fade_start) * 0.02)  # Decrease alpha
            selected_color = random.choice(color_palette)
            # Convert color to RGBA
            selected_color_rgb = mcolors.to_rgb(selected_color)
            face_color = selected_color_rgb + (alpha,)
        else:
            alpha = 0.9
            selected_color = random.choice(color_palette)
            face_color = selected_color  # RGBA not needed as alpha=0.9 is already set

        # Create the cube
        if i >= fade_start:
            poly3d = Poly3DCollection(vertices, facecolors=face_color, linewidths=0.5, edgecolors='k', alpha=alpha)
        else:
            poly3d = Poly3DCollection(vertices, facecolors=face_color, linewidths=0.5, edgecolors='k', alpha=0.9)
        ax.add_collection3d(poly3d)

        # Add ellipsis at the end to symbolize infinity
        if i == num_cubes -1:
            ax.text(r[0] + cube_size / 2, r[1] + cube_size / 2, r[2] + cube_size / 2, "...", 
                    color='black', ha='center', va='center', fontsize=12, weight='bold')

    # Set the axes limits
    ax.set_xlim(-2, num_cubes * spacing * 1.2 + 2)
    ax.set_ylim(-2, num_cubes * spacing * 0.6 + 4)
    ax.set_zlim(-1, 2)

    # Remove axis ticks and panes with updated attribute references
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    ax.xaxis.pane.fill = False  # Updated from w_xaxis
    ax.yaxis.pane.fill = False  # Updated from w_yaxis
    ax.zaxis.pane.fill = False  # Updated from w_zaxis

    # Set viewing angle for better perception
    ax.view_init(elev=30, azim=45)

    # Set title
    plt.title("Infinity Array Visualization", fontsize=22, pad=20)

    # Apply a light grid to enhance depth perception
    ax.grid(False)

    # Increase background color for better contrast
    ax.set_facecolor('white')

    # Adjust layout
    plt.tight_layout()

    # Save the image
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"Infinity array image saved to {save_path}")

    # Clear the figure to free memory
    plt.clf()

def generate_random_array(size, min_val=0, max_val=100):
    """
    Generates a random array of a given size with integer values between min_val and max_val.
    """
    return [random.randint(min_val, max_val) for _ in range(size)]

def create_dataset(output_dir='output_arrays', max_size=25, samples_per_size=100):
    """
    Creates a dataset of array images from size 0 to max_size and an infinity array.
    Organized into separate directories for each array size with varied color schemes.
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    metadata = []
    print(f"Creating array images in directory: {output_dir}")

    # Define a list of color palettes to introduce variety
    color_palettes = [
        ['skyblue', 'lightgreen', 'salmon', 'plum', 'lightcoral'],
        ['lightseagreen', 'paleturquoise', 'lightpink', 'lightsalmon', 'mediumslateblue'],
        ['lightskyblue', 'lightgoldenrodyellow', 'lightsteelblue', 'mistyrose', 'lavender'],
        ['khaki', 'lightcyan', 'thistle', 'peachpuff', 'wheat']
    ]

    num_palettes = len(color_palettes)

    # Generate arrays from size 0 to max_size with multiple samples
    for size in range(0, max_size + 1):
        # Create a directory for the current array size
        size_dir = os.path.join(output_dir, str(size))
        if not os.path.exists(size_dir):
            os.makedirs(size_dir)

        for sample in range(1, samples_per_size + 1):
            if size == 0:
                array = []
                title = f"Array Size: {size} (Empty Array)"
                filename = f"array_size_{size}_sample_{sample}_empty.png"
            else:
                array = generate_random_array(size)
                title = f"Array Size: {size}"
                filename = f"array_size_{size}_sample_{sample}.png"

            # Define save path
            save_path = os.path.join(size_dir, filename)

            # Select a color palette in a round-robin fashion
            palette_idx = (sample - 1) % num_palettes
            selected_palette = color_palettes[palette_idx]

            # Assign a random color from the selected palette to each cube
            # This ensures varied color schemes across samples
            cube_colors = [random.choice(selected_palette) for _ in range(size)]

            if size == 0:
                # Draw the empty array
                draw_3d_array(
                    array_values=array,
                    title=title,
                    index_color='black',
                    value_color='black',
                    show_labels=True,
                    elevation=random.randint(20, 40),  # Randomize elevation for diversity
                    azimuth=random.randint(30, 60),    # Randomize azimuth for diversity
                    save_path=save_path
                )
            else:
                # Draw and save the image with varied colors
                fig = plt.figure(figsize=(8, 6))
                ax = fig.add_subplot(111, projection='3d')

                # Parameters
                cube_size = 0.8
                spacing = 1.0

                for i, val in enumerate(array):
                    # Define the starting corner of the cube
                    r = [i * spacing, 0, 0]
                    # Apply slight random variation
                    r_variation = [random.uniform(-0.05, 0.05) for _ in range(3)]
                    r = [r[j] + r_variation[j] for j in range(3)]
                    points = np.array([
                        [r[0], r[1], r[2]],
                        [r[0] + cube_size, r[1], r[2]],
                        [r[0] + cube_size, r[1] + cube_size, r[2]],
                        [r[0], r[1] + cube_size, r[2]],
                        [r[0], r[1], r[2] + cube_size],
                        [r[0] + cube_size, r[1], r[2] + cube_size],
                        [r[0] + cube_size, r[1] + cube_size, r[2] + cube_size],
                        [r[0], r[1] + cube_size, r[2] + cube_size]
                    ])

                    # Define the vertices for each face of the cube
                    vertices = [
                        [points[0], points[1], points[2], points[3]],  # Bottom face
                        [points[4], points[5], points[6], points[7]],  # Top face
                        [points[0], points[1], points[5], points[4]],  # Front face
                        [points[2], points[3], points[7], points[6]],  # Back face
                        [points[1], points[2], points[6], points[5]],  # Right face
                        [points[4], points[7], points[3], points[0]]   # Left face
                    ]

                    # Assign the pre-selected cube color
                    selected_color = cube_colors[i]
                    # Convert color to RGBA if needed for alpha
                    if isinstance(selected_color, str):
                        selected_color_rgb = mcolors.to_rgb(selected_color)
                    else:
                        selected_color_rgb = selected_color
                    face_color = selected_color_rgb + (0.9,)  # Adding alpha for slight transparency

                    # Create the cube
                    poly3d = Poly3DCollection(vertices, facecolors=face_color, linewidths=0.5, edgecolors='k', alpha=0.9)
                    ax.add_collection3d(poly3d)

                    # Add index and value labels with slight randomization in position
                    if True:
                        # Index label below the cube
                        ax.text(r[0] + cube_size / 2, r[1] - 0.3, r[2], f"{i}", color='black',
                                ha='center', va='top', fontsize=8, weight='bold')
                        # Value label inside the cube
                        ax.text(r[0] + cube_size / 2, r[1] + cube_size / 2, r[2] + cube_size / 2, str(val), 
                                color='white', ha='center', va='center', fontsize=10, weight='bold')

                # Set the axes limits
                ax.set_xlim(-1, len(array) * spacing)
                ax.set_ylim(-1, 2)
                ax.set_zlim(-1, 2)

                # Remove axis ticks and panes with updated attribute references
                ax.set_xticks([])
                ax.set_yticks([])
                ax.set_zticks([])
                ax.xaxis.pane.fill = False  # Updated from w_xaxis
                ax.yaxis.pane.fill = False  # Updated from w_yaxis
                ax.zaxis.pane.fill = False  # Updated from w_zaxis

                # Set viewing angle
                ax.view_init(elev=random.randint(20, 40), azim=random.randint(30, 60))

                # Set title
                plt.title(title, fontsize=16, pad=20)

                # Set background color
                ax.set_facecolor('white')

                # Adjust layout
                plt.tight_layout()

                # Save the image with varied colors
                plt.savefig(save_path, dpi=300, bbox_inches='tight')
                plt.clf()

            # Add to metadata
            metadata.append({
                'size': size,
                'elements': array,
                'image': save_path
            })

    print(f"Generated {samples_per_size} samples for array size {size}.")

    # Generate the infinity array once
    infinity_save_path = os.path.join(output_dir, "infinity_array.png")
    draw_infinity_array(save_path=infinity_save_path)
    metadata.append({
        'size': 'infty',
        'elements': 'infinite',
        'image': infinity_save_path
    })
    print("Generated infinity array visualization.")

    # Save metadata to a JSON file
    metadata_path = os.path.join(output_dir, 'metadata.json')
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=4)
    print(f"Metadata saved to {metadata_path}")

    print("Dataset creation complete.")

if __name__ == "__main__":
    # Parameters
    OUTPUT_DIR = 'output_arrays'    # Directory to save images
    MAX_ARRAY_SIZE = 25             # Maximum array size
    SAMPLES_PER_SIZE = 100          # Number of images per array size

    create_dataset(output_dir=OUTPUT_DIR, max_size=MAX_ARRAY_SIZE, samples_per_size=SAMPLES_PER_SIZE)

this is how i made the dataset

then i cgan comes
import os
import json
import random
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from PIL import Image
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.colors as mcolors
from tqdm import tqdm

class Generator(nn.Module):
    def __init__(self, input_dim=25, noise_dim=100, feature_g=64, img_channels=3):
        """
        Initializes the Generator network.
        
        Parameters:
        - input_dim: Dimension of the encoded array input.
        - noise_dim: Dimension of the noise vector.
        - feature_g: Base number of features in generator.
        - img_channels: Number of image channels (e.g., 3 for RGB).
        """
        super(Generator, self).__init__()
        self.noise_dim = noise_dim
        self.net = nn.Sequential(
            # Input: (input_dim + noise_dim) x 1 x 1
            nn.ConvTranspose2d(input_dim + noise_dim, feature_g * 8, kernel_size=4, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(feature_g * 8),
            nn.ReLU(True),
            # State: (feature_g*8) x 4 x 4
            nn.ConvTranspose2d(feature_g * 8, feature_g * 4, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(feature_g * 4),
            nn.ReLU(True),
            # State: (feature_g*4) x 8 x 8
            nn.ConvTranspose2d(feature_g * 4, feature_g * 2, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(feature_g * 2),
            nn.ReLU(True),
            # State: (feature_g*2) x 16 x 16
            nn.ConvTranspose2d(feature_g * 2, feature_g, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(feature_g),
            nn.ReLU(True),
            # State: (feature_g) x 32 x 32
            nn.ConvTranspose2d(feature_g, img_channels, kernel_size=4, stride=2, padding=1, bias=False),
            nn.Tanh()
            # Output: (img_channels) x 64 x 64
        )
    
    def forward(self, x, noise):
        """
        Forward pass of the Generator.
        
        Parameters:
        - x: Encoded array tensor.
        - noise: Random noise tensor.
        
        Returns:
        - Generated image tensor.
        """
        # Concatenate array encoding and noise
        combined = torch.cat([x, noise], dim=1)
        return self.net(combined)
class Discriminator(nn.Module):
    def __init__(self, img_channels=3, feature_d=64, input_dim=25):
        """
        Initializes the Discriminator network.
        
        Parameters:
        - img_channels: Number of image channels (e.g., 3 for RGB).
        - feature_d: Base number of features in discriminator.
        - input_dim: Dimension of the encoded array input.
        """
        super(Discriminator, self).__init__()
        # Embedding the array condition
        self.label_embedding = nn.Sequential(
            nn.Linear(input_dim, input_dim),
            nn.ReLU(True)
        )
        self.net = nn.Sequential(
            # Input: (img_channels + input_dim) x 64 x 64
            nn.Conv2d(img_channels + input_dim, feature_d, kernel_size=4, stride=2, padding=1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # State: (feature_d) x 32 x 32
            nn.Conv2d(feature_d, feature_d * 2, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(feature_d * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # State: (feature_d*2) x 16 x 16
            nn.Conv2d(feature_d * 2, feature_d * 4, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(feature_d * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # State: (feature_d*4) x 8 x 8
            nn.Conv2d(feature_d * 4, feature_d * 8, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(feature_d * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # State: (feature_d*8) x 4 x 4
            nn.Conv2d(feature_d * 8, 1, kernel_size=4, stride=1, padding=0, bias=False),
            nn.Sigmoid()
            # Output: 1
        )
    
    def forward(self, img, labels):
        """
        Forward pass of the Discriminator.
        
        Parameters:
        - img: Image tensor.
        - labels: Encoded array tensor.
        
        Returns:
        - Probability tensor indicating real or fake.
        """
        # Embed the labels
        label_embedding = self.label_embedding(labels).unsqueeze(2).unsqueeze(3)  # Shape: (batch_size, input_dim, 1, 1)
        label_embedding = label_embedding.repeat(1, 1, img.size(2), img.size(3))  # Shape: (batch_size, input_dim, H, W)
        # Concatenate image and label embedding
        combined = torch.cat([img, label_embedding], dim=1)  # Shape: (batch_size, img_channels + input_dim, H, W)
        return self.net(combined).view(-1)
class ArrayImageDataset(Dataset):
    def __init__(self, metadata_path, transform=None, max_size=25):
        """
        Initializes the dataset by loading metadata and setting up transforms.
        
        Parameters:
        - metadata_path: Path to the metadata JSON file.
        - transform: Optional torchvision transforms to apply to images.
        - max_size: Maximum size of arrays for encoding.
        """
        with open(metadata_path, 'r') as f:
            self.metadata = json.load(f)
        self.transform = transform
        self.max_size = max_size

    def __len__(self):
        return len(self.metadata)

    def encode_array(self, array):
        """
        Encodes the array by padding it to `max_size` with -1.
        
        Parameters:
        - array: List of integers representing the array.
        
        Returns:
        - Padded numpy array.
        """
        padded = array + [-1] * (self.max_size - len(array))
        return padded

    def __getitem__(self, idx):
        """
        Retrieves the array data and corresponding image.
        
        Parameters:
        - idx: Index of the data point.
        
        Returns:
        - input_vector: Tensor of encoded array data.
        - image: Transformed image tensor.
        """
        data_point = self.metadata[idx]
        array = data_point['elements']
        image_path = data_point['image']
        size = data_point['size']

        # Encode the array
        encoded_array = self.encode_array(array)
        encoded_array = np.array(encoded_array, dtype=np.float32)

        # Convert to tensor
        input_vector = torch.from_numpy(encoded_array)

        # Load and transform the image
        image = Image.open(image_path).convert('RGB')
        if self.transform:
            image = self.transform(image)

        return input_vector, image
transform = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)  # Normalize to [-1, 1]
])
metadata_path = '/kaggle/working/output_arrays/metadata.json'  # Update this path if different
data_dir = '/kaggle/working/output_arrays'

dataset = ArrayImageDataset(metadata_path=metadata_path, transform=transform, max_size=25)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=4)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Initialize Generator and Discriminator
generator = Generator(input_dim=25, noise_dim=100).to(device)
discriminator = Discriminator(img_channels=3, feature_d=64, input_dim=25).to(device)

# Loss function
criterion = nn.BCELoss()

# Optimizers
optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
# Function to save generated images for monitoring progress
def save_generated_images(epoch, generator, fixed_input, noise_dim=100, output_dir='generated_images'):
    """
    Generates and saves images from the fixed input arrays to monitor training progress.
    
    Parameters:
    - epoch: Current epoch number.
    - generator: Trained Generator model.
    - fixed_input: Fixed input tensor.
    - noise_dim: Dimension of the noise vector.
    - output_dir: Directory to save generated images.
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    with torch.no_grad():
        noise = torch.randn(fixed_input.size(0), noise_dim, 1, 1, device=device)
        fake_images = generator(fixed_input, noise).detach().cpu()
    grid = utils.make_grid(fake_images, nrow=5, normalize=True)
    save_image(grid, os.path.join(output_dir, f'epoch_{epoch}.png'))
    print(f"Saved generated images for epoch {epoch}.")
# Define fixed sample arrays
sample_arrays = [
    [],  # Empty array
    [10],  # Single-element array
    [5, 3, 8, 6, 2],  # Five-element array
    [7, 9, 4, 1, 3, 8, 2],  # Seven-element array
    [15, 23, 7, 19, 12, 16, 9, 4]  # Eight-element array
]

# Function to encode these fixed arrays
def encode_fixed_arrays(arrays, max_size=25):
    """
    Encodes multiple arrays into a tensor suitable for the generator.
    """
    encoded = []
    for array in arrays:
        padded = array + [-1]*(max_size - len(array))
        encoded.append(padded)
    tensor = torch.tensor(encoded, dtype=torch.float32)
    tensor = tensor.unsqueeze(2).unsqueeze(3)  # Shape: (num_samples, max_size, 1, 1)
    return tensor.to(device)

fixed_input = encode_fixed_arrays(sample_arrays, max_size=25)
num_epochs = 100  # Adjust based on convergence
noise_dim = 100  # Dimension of the noise vector

for epoch in range(1, num_epochs + 1):
    generator.train()
    discriminator.train()
    loop = tqdm(dataloader, desc=f'Epoch [{epoch}/{num_epochs}]', leave=False)
    
    for batch_idx, (inputs, real_images) in enumerate(loop):
        batch_size_curr = real_images.size(0)
        inputs = inputs.to(device)
        real_images = real_images.to(device)
        
        # ========================================
        # Train Discriminator
        # ========================================
        discriminator.zero_grad()
        
        # Train on real images
        labels_real = torch.full((batch_size_curr,), 0.9, dtype=torch.float, device=device)  # Label smoothing
        output_real = discriminator(real_images, inputs).view(-1)
        loss_D_real = criterion(output_real, labels_real)
        loss_D_real.backward()
        
        # Generate fake images
        noise = torch.randn(batch_size_curr, noise_dim, 1, 1, device=device)
        fake_images = generator(inputs, noise)
        labels_fake = torch.full((batch_size_curr,), 0.0, dtype=torch.float, device=device)
        output_fake = discriminator(fake_images.detach(), inputs).view(-1)
        loss_D_fake = criterion(output_fake, labels_fake)
        loss_D_fake.backward()
        
        loss_D = loss_D_real + loss_D_fake
        optimizer_D.step()
        
        # ========================================
        # Train Generator
        # ========================================
        generator.zero_grad()
        labels_gen = torch.full((batch_size_curr,), 0.9, dtype=torch.float, device=device)  # Generator aims to fool discriminator
        output_gen = discriminator(fake_images, inputs).view(-1)
        loss_G = criterion(output_gen, labels_gen)
        loss_G.backward()
        optimizer_G.step()
        
        # Update tqdm loop description
        loop.set_postfix(loss_D=loss_D.item(), loss_G=loss_G.item())
    
    # Save generated images at specified intervals
    if epoch % 10 == 0 or epoch == 1:
        save_generated_images(epoch, generator, fixed_input, noise_dim=noise_dim)
# Directory to save models
model_save_dir = 'trained_models'
os.makedirs(model_save_dir, exist_ok=True)

# Save the Generator and Discriminator state dictionaries
torch.save(generator.state_dict(), os.path.join(model_save_dir, 'generator.pth'))
torch.save(discriminator.state_dict(), os.path.join(model_save_dir, 'discriminator.pth'))
print(f"Generator and Discriminator models saved to '{model_save_dir}' directory.")

do you find any specific mistake in my code because this came up
                                                   
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[18], line 27
     25 # Generate fake images
     26 noise = torch.randn(batch_size_curr, noise_dim, 1, 1, device=device)
---> 27 fake_images = generator(inputs, noise)
     28 labels_fake = torch.full((batch_size_curr,), 0.0, dtype=torch.float, device=device)
     29 output_fake = discriminator(fake_images.detach(), inputs).view(-1)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

Cell In[10], line 49, in Generator.forward(self, x, noise)
     38 """
     39 Forward pass of the Generator.
     40 
   (...)
     46 - Generated image tensor.
     47 """
     48 # Concatenate array encoding and noise
---> 49 combined = torch.cat([x, noise], dim=1)
     50 return self.net(combined)

RuntimeError: Tensors must have same number of dimensions: got 2 and 4


also i would like to make some new changes, so i want to save checkpoints maybe so if training stops somwhere it cntinues fro there after getting over with the error only, and it shows the validation result and stuff after every epoch so uk ik how is things going