2025 USA-NA-AIO Round 2, Problem 1, Part 4

AI Olympiads 2025 USA-NA-AIO Round 2
1

Part 4 (10 points, coding task)

In this part, you are asked to create the dataset \mathcal D_{PDE}.

  1. The dataset object is called dataset_train_PDE. It is in a class called Dataset_PDE that you need to build.

  2. Class Dataset_PDE subclasses Dataset.

  3. Each \left( t, x \right) \in \mathcal D_{PDE} is randomly sampled from \left[ 0 , 1 \right]^2.

  4. Set | \mathcal D_{PDE} | = 500.

2 
### WRITE YOUR SOLUTION HERE ###

num_samples = 500

class Dataset_PDE(Dataset):
    def __init__(self, num_samples):
        self.data = torch.rand(num_samples, 2, requires_grad=True)

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, idx):
        return self.data[idx]

dataset_train_PDE = Dataset_PDE(num_samples)

""" END OF THIS PART """