code: def make_dataset(t0, t1, batch_size, noise_std, train_dir, device):
      data_path = os.path.join(train_dir, "lorenz_data.pth")
      if os.path.exists(data_path):
          data_dict = torch.load(data_path)
          xs, ts = data_dict["xs"], data_dict["ts"]
          logging.warning(f"Loaded toy data at: {data_path}")
          if xs.shape[1] != batch_size:
              raise ValueError(
                  "Batch size has changed; please delete and regenerate the data."
              )
          if ts[0] != t0 or ts[-1] != t1:
              raise ValueError(
                  "Times interval [t0, t1] has changed; please delete and regenerate the data."
              )
      else:
          _y0 = torch.randn(batch_size, 3, device=device)
          ts = torch.linspace(t0, t1, steps=100, device=device)
          xs = StochasticLorenz().sample(
              _y0, ts, noise_std, normalize=True
          )
  
          os.makedirs(os.path.dirname(data_path), exist_ok=True)
          torch.save({"xs": xs, "ts": ts}, data_path)
          logging.warning(f"Stored toy data at: {data_path}")
      return xs, ts
  
imports: import logging
  import os
  from typing import Sequence
  import fire
  import matplotlib.gridspec
  import matplotlib.pyplot
  import numpy
  import torch
  import tqdm
  from torch import nn
  from torch import optim
  from torch.distributions import Normal
  import torchsde
label_desc: Use image node_id 1 for calling a global function with example usage: make_dataset(t0, t1, batch_size, noise_std, train_dir, device) and returns: xs,  ts