code: def get_data(train_size=1000, valid_size=400):
      train_x, train_y = get_linear_data(2, 5, train_size)
      val_x, val_y = get_linear_data(2, 5, valid_size)
      data = {
          "x": train_x,
          "y": train_y,
          "val_x": val_x,
          "val_y": val_y,
      }
      return data
  
imports: from unittest import TestCase
  from zoo.orca.automl.model.base_pytorch_model import PytorchModelBuilder
  from torch.utils.data import Dataset, DataLoader
  import pytest
  import numpy
  import torch
  import torch.nn
label_desc: Use image node_id 2 for calling a global function with example usage: get_data(train_size, valid_size) and returns: data