code: def train_data_creator(a=5, b=10, size=1000):
      def get_dataset(a, b, size) -> Dataset:
          items = [i / size for i in range(size)]
          dataset = ray.data.from_items(
              [{"x": x, "y": a * x + b} for x in items]
          )
          return dataset
  
      train_dataset = get_dataset(a, b, size)
      return train_dataset
  
imports: import pytest
  import logging
  from unittest import TestCase
  import ray
  from ray.data import Dataset
  import torch.nn
  import torch.optim
  from bigdl.orca import init_orca_context, stop_orca_context
  from bigdl.orca.learn.pytorch import Estimator
  from bigdl.orca.learn.metrics import MSE
label_desc: Use image node_id 1 for calling a global function with example usage: train_data_creator(a, b, size) and returns: train_dataset, dataset