code: def get_data(batch_size, device):
      dataset_size = 8192
      t_size = 64
  
      class OrnsteinUhlenbeckSDE(torch.nn.Module):
          sde_type = "ito"
          noise_type = "scalar"
  
          def __init__(self, mu, theta, sigma):
              super().__init__()
              self.register_buffer("mu", torch.as_tensor(mu))
              self.register_buffer("theta", torch.as_tensor(theta))
              self.register_buffer("sigma", torch.as_tensor(sigma))
  
          def f(self, t, y):
              return self.mu * t - self.theta * y
  
          def g(self, t, y):
              return self.sigma.expand(y.size(0), 1, 1) * (
                  2 * t / t_size
              )
  
      ou_sde = OrnsteinUhlenbeckSDE(mu=0.02, theta=0.1, sigma=0.4).to(
          device
      )
      y0 = torch.rand(dataset_size, device=device).unsqueeze(-1) * 2 - 1
      ts = torch.linspace(0, t_size - 1, t_size, device=device)
      ys = torchsde.sdeint(ou_sde, y0, ts, dt=1e-1)
  
      ###################
      # To demonstrate how to handle irregular data, then here we additionally drop some of the data (by setting it to
      # NaN.)
      ###################
      ys_num = ys.numel()
      to_drop = torch.randperm(ys_num)[: int(0.3 * ys_num)]
      ys.view(-1)[to_drop] = float("nan")
  
      ###################
      # Typically important to normalise data. Note that the data is normalised with respect to the statistics of the
      # initial data, _not_ the whole time series. This seems to help the learning process, presumably because if the
      # initial condition is wrong then it's pretty hard to learn the rest of the SDE correctly.
      ###################
      y0_flat = ys[0].view(-1)
      y0_not_nan = y0_flat.masked_select(~torch.isnan(y0_flat))
      ys = (ys - y0_not_nan.mean()) / y0_not_nan.std()
  
      ###################
      # As discussed, time must be included as a channel for the discriminator.
      ###################
      ys = torch.cat(
          [
              ts.unsqueeze(0)
              .unsqueeze(-1)
              .expand(dataset_size, t_size, 1),
              ys.transpose(0, 1),
          ],
          dim=2,
      )
      # shape (dataset_size=1000, t_size=100, 1 + data_size=3)
  
      ###################
      # Package up.
      ###################
      data_size = (
          ys.size(-1) - 1
      )  # How many channels the data has (not including time, hence the minus one).
      ys_coeffs = torchcde.linear_interpolation_coeffs(
          ys
      )  # as per neural CDEs.
      dataset = torch.utils.data.TensorDataset(ys_coeffs)
      dataloader = torch.utils.data.DataLoader(
          dataset, batch_size=batch_size, shuffle=True
      )
  
      return ts, data_size, dataloader
  
imports: import fire
  import matplotlib.pyplot
  import torch
  import torch.optim.swa_utils
  import torchcde
  import torchsde
  import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: get_data(batch_size, device) and returns: ts,  data_size,  dataloader, unknown, unknown