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

# ========================================
# Section 1: Importing Necessary Libraries
# ========================================
import os
import json
import random
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import transforms, utils
from torchvision.utils import save_image
from tqdm import tqdm

# ========================================
# Section 2: Defining the Dataset Class
# ========================================
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.
        """
        if not os.path.exists(metadata_path):
            raise FileNotFoundError(f"Metadata file not found at: {metadata_path}")
        
        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 with shape (max_size, 1, 1).
        - image: Transformed image tensor.
        """
        data_point = self.metadata[idx]
        array = data_point['elements']
        image_path = data_point['image']
        size = data_point['size']

        # Handle infinity arrays or special cases
        if isinstance(array, str) and array.lower() == 'infinite':
            array = []  # Representing as an empty array or any other desired encoding

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

        # Convert to tensor and reshape to (max_size, 1, 1)
        input_vector = torch.from_numpy(encoded_array).unsqueeze(1).unsqueeze(2)  # Shape: (max_size, 1, 1)

        # Load and transform the image
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image not found at: {image_path}")
        
        image = Image.open(image_path).convert('RGB')
        if self.transform:
            image = self.transform(image)

        return input_vector, image

# ========================================
# Section 3: Building the Generator and Discriminator Models
# ========================================
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 of shape (batch_size, input_dim, 1, 1)
        - noise: Random noise tensor of shape (batch_size, noise_dim, 1, 1)
        
        Returns:
        - Generated image tensor.
        """
        # Concatenate array encoding and noise along the channel dimension
        combined = torch.cat([x, noise], dim=1)  # Shape: (batch_size, input_dim + noise_dim, 1, 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 of shape (batch_size, img_channels, 64, 64)
        - labels: Encoded array tensor of shape (batch_size, input_dim, 1, 1)
        
        Returns:
        - Probability tensor indicating real or fake.
        """
        batch_size = img.size(0)
        
        # Embed the labels
        labels = labels.view(batch_size, -1)  # Shape: (batch_size, input_dim)
        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)

# ========================================
# Section 4: Utility Functions
# ========================================
def save_checkpoint(state, filename='checkpoint.pth.tar'):
    """
    Saves the training state to a checkpoint file.
    
    Parameters:
    - state: Dictionary containing model and optimizer states.
    - filename: Name of the checkpoint file.
    """
    torch.save(state, filename)
    print(f"Checkpoint saved to {filename}")

def load_checkpoint(filename, generator, discriminator, optimizer_G, optimizer_D, device):
    """
    Loads the training state from a checkpoint file.
    
    Parameters:
    - filename: Path to the checkpoint file.
    - generator: Generator model instance.
    - discriminator: Discriminator model instance.
    - optimizer_G: Generator optimizer instance.
    - optimizer_D: Discriminator optimizer instance.
    - device: Device to map the checkpoint.
    
    Returns:
    - epoch: Resumed epoch number.
    - loss_D: Last recorded Discriminator loss.
    - loss_G: Last recorded Generator loss.
    """
    if os.path.isfile(filename):
        print(f"Loading checkpoint '{filename}'")
        checkpoint = torch.load(filename, map_location=device)
        generator.load_state_dict(checkpoint['generator_state_dict'])
        discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
        optimizer_G.load_state_dict(checkpoint['optimizer_G_state_dict'])
        optimizer_D.load_state_dict(checkpoint['optimizer_D_state_dict'])
        epoch = checkpoint['epoch']
        loss_D = checkpoint.get('loss_D', None)
        loss_G = checkpoint.get('loss_G', None)
        print(f"Loaded checkpoint '{filename}' (epoch {epoch})")
        return epoch, loss_D, loss_G
    else:
        print(f"No checkpoint found at '{filename}'")
        return 1, None, None

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}.")

# ========================================
# Section 5: Training Setup and Execution
# ========================================
def main():
    # Define transformations for the images
    transform = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)  # Normalize to [-1, 1]
    ])
    
    # Paths
    metadata_path = '/kaggle/input/cgandb/metadata.json'  # Update this path if different
    data_dir = '/kaggle/input/cgandb'  # Directory containing images and metadata
    
    # Initialize the dataset
    dataset = ArrayImageDataset(metadata_path=metadata_path, transform=transform, max_size=25)
    
    # Split into training and validation sets (80% train, 20% val)
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
    
    # Create DataLoaders
    train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
    val_dataloader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)
    
    # Device configuration
    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))
    
    # Checkpoint configuration
    CHECKPOINT_DIR = 'checkpoints'
    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
    CHECKPOINT_PATH = os.path.join(CHECKPOINT_DIR, 'checkpoint.pth.tar')
    
    # Load checkpoint if it exists
    start_epoch, _, _ = load_checkpoint(CHECKPOINT_PATH, generator, discriminator, optimizer_G, optimizer_D, device)
    
    # Define fixed sample arrays for monitoring
    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
    ]
    
    def encode_fixed_arrays(arrays, max_size=25):
        """
        Encodes multiple arrays into a tensor suitable for the generator.
        
        Parameters:
        - arrays: List of arrays to encode.
        - max_size: Maximum size of arrays.
        
        Returns:
        - Tensor of shape (num_samples, max_size, 1, 1)
        """
        encoded = []
        for array in arrays:
            if isinstance(array, str) and array.lower() == 'infinite':
                array = []  # Adjust as needed
            padded = array + [-1]*(max_size - len(array))
            encoded.append(padded)
        tensor = torch.tensor(encoded, dtype=torch.float32)
        tensor = tensor.unsqueeze(1).unsqueeze(2)  # Shape: (num_samples, max_size, 1, 1)
        return tensor.to(device)
    
    # Encode fixed input arrays
    fixed_input = encode_fixed_arrays(sample_arrays, max_size=25)
    noise_dim = 100
    num_epochs = 100  # Adjust based on convergence
    
    # Training Loop
    for epoch in range(start_epoch, num_epochs + 1):
        generator.train()
        discriminator.train()
        loop = tqdm(train_dataloader, desc=f'Epoch [{epoch}/{num_epochs}]', leave=False)
        total_loss_D = 0.0
        total_loss_G = 0.0
        
        for batch_idx, (inputs, real_images) in enumerate(loop):
            batch_size_curr = real_images.size(0)
            inputs = inputs.to(device)  # Shape: (batch_size, max_size, 1, 1)
            real_images = real_images.to(device)
            
            # ========================================
            # Train Discriminator
            # ========================================
            discriminator.zero_grad()
            
            # 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()
            
            # 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()
            
            total_loss_D += loss_D.item()
            total_loss_G += loss_G.item()
            
            # Update tqdm loop description
            loop.set_postfix(loss_D=loss_D.item(), loss_G=loss_G.item())
        
        # Calculate average losses
        avg_loss_D = total_loss_D / len(train_dataloader)
        avg_loss_G = total_loss_G / len(train_dataloader)
        
        # ========================================
        # Validation Phase
        # ========================================
        generator.eval()
        discriminator.eval()
        val_loss_D = 0.0
        val_loss_G = 0.0
        with torch.no_grad():
            for inputs, real_images in val_dataloader:
                batch_size_curr = real_images.size(0)
                inputs = inputs.to(device)
                real_images = real_images.to(device)
                
                # Real images
                labels_real = torch.full((batch_size_curr,), 0.9, dtype=torch.float, device=device)
                output_real = discriminator(real_images, inputs).view(-1)
                loss_D_real = criterion(output_real, labels_real)
                
                # 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, inputs).view(-1)
                loss_D_fake = criterion(output_fake, labels_fake)
                
                loss_D = loss_D_real + loss_D_fake
                
                # Generator loss on validation
                labels_gen = torch.full((batch_size_curr,), 0.9, dtype=torch.float, device=device)
                output_gen = discriminator(fake_images, inputs).view(-1)
                loss_G = criterion(output_gen, labels_gen)
                
                val_loss_D += loss_D.item()
                val_loss_G += loss_G.item()
        
        # Compute average validation losses
        val_loss_D /= len(val_dataloader)
        val_loss_G /= len(val_dataloader)
        
        # Print training and validation results
        print(f"Epoch [{epoch}/{num_epochs}] | "
              f"Train Loss_D: {avg_loss_D:.4f}, Loss_G: {avg_loss_G:.4f} | "
              f"Val Loss_D: {val_loss_D:.4f}, Loss_G: {val_loss_G:.4f}")
        
        # Save generated images at specified intervals
        if epoch % 10 == 0 or epoch == 1:
            save_generated_images(epoch, generator, fixed_input, noise_dim=noise_dim)
        
        # Save checkpoint at the end of each epoch
        checkpoint_state = {
            'epoch': epoch + 1,
            'generator_state_dict': generator.state_dict(),
            'discriminator_state_dict': discriminator.state_dict(),
            'optimizer_G_state_dict': optimizer_G.state_dict(),
            'optimizer_D_state_dict': optimizer_D.state_dict(),
            'loss_D': avg_loss_D,
            'loss_G': avg_loss_G,
        }
        save_checkpoint(checkpoint_state, filename=CHECKPOINT_PATH)
    
    # ========================================
    # Section 6: Saving the Final Models
    # ========================================
    model_save_dir = 'trained_models'
    os.makedirs(model_save_dir, exist_ok=True)
    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.")

# ========================================
# Section 7: Encoding Fixed Input Arrays
# ========================================
def encode_fixed_arrays(arrays, max_size=25, device='cpu'):
    """
    Encodes multiple arrays into a tensor suitable for the generator.
    
    Parameters:
    - arrays: List of arrays to encode.
    - max_size: Maximum size of arrays.
    - device: Device to place the tensor on.
    
    Returns:
    - Tensor of shape (num_samples, max_size, 1, 1)
    """
    encoded = []
    for array in arrays:
        if isinstance(array, str) and array.lower() == 'infinite':
            array = []  # Adjust as needed
        padded = array + [-1]*(max_size - len(array))
        encoded.append(padded)
    tensor = torch.tensor(encoded, dtype=torch.float32)
    tensor = tensor.unsqueeze(1).unsqueeze(2)  # Shape: (num_samples, max_size, 1, 1)
    return tensor.to(device)

# ========================================
# Section 8: Running the Script
# ========================================
if __name__ == "__main__":
    main()


but now the issue
Using device: cuda
No checkpoint found at 'checkpoints/checkpoint.pth.tar'
                                                     
---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
Cell In[2], line 514
    510 # ========================================
    511 # Section 8: Running the Script
    512 # ========================================
    513 if __name__ == "__main__":
--> 514     main()

Cell In[2], line 369, in main()
    366 total_loss_D = 0.0
    367 total_loss_G = 0.0
--> 369 for batch_idx, (inputs, real_images) in enumerate(loop):
    370     batch_size_curr = real_images.size(0)
    371     inputs = inputs.to(device)  # Shape: (batch_size, max_size, 1, 1)

File /opt/conda/lib/python3.10/site-packages/tqdm/std.py:1181, in tqdm.__iter__(self)
   1178 time = self._time
   1180 try:
-> 1181     for obj in iterable:
   1182         yield obj
   1183         # Update and possibly print the progressbar.
   1184         # Note: does not call self.update(1) for speed optimisation.

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.__next__(self)
    627 if self._sampler_iter is None:
    628     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    629     self._reset()  # type: ignore[call-arg]
--> 630 data = self._next_data()
    631 self._num_yielded += 1
    632 if self._dataset_kind == _DatasetKind.Iterable and \
    633         self._IterableDataset_len_called is not None and \
    634         self._num_yielded > self._IterableDataset_len_called:

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1344, in _MultiProcessingDataLoaderIter._next_data(self)
   1342 else:
   1343     del self._task_info[idx]
-> 1344     return self._process_data(data)

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1370, in _MultiProcessingDataLoaderIter._process_data(self, data)
   1368 self._try_put_index()
   1369 if isinstance(data, ExceptionWrapper):
-> 1370     data.reraise()
   1371 return data

File /opt/conda/lib/python3.10/site-packages/torch/_utils.py:706, in ExceptionWrapper.reraise(self)
    702 except TypeError:
    703     # If the exception takes multiple arguments, don't try to
    704     # instantiate since we don't know how to
    705     raise RuntimeError(msg) from None
--> 706 raise exception

FileNotFoundError: Caught FileNotFoundError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 309, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 50, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/dataset.py", line 420, in __getitems__
    return [self.dataset[self.indices[idx]] for idx in indices]
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/data/dataset.py", line 420, in <listcomp>
    return [self.dataset[self.indices[idx]] for idx in indices]
  File "/tmp/ipykernel_30/3041678846.py", line 87, in __getitem__
    raise FileNotFoundError(f"Image not found at: {image_path}")
FileNotFoundError: Image not found at: output_arrays/13/array_size_13_sample_70.png


my current working directory
import os
print(os.getcwd())
/kaggle/working


so please tell me where did i go rong?