code: def plot(ts, generator, dataloader, num_plot_samples, plot_locs):
      # Get samples
      (real_samples,) = next(iter(dataloader))
      assert num_plot_samples <= real_samples.size(0)
      real_samples = torchcde.LinearInterpolation(
          real_samples
      ).evaluate(ts)
      real_samples = real_samples[..., 1]
  
      with torch.no_grad():
          generated_samples = generator(ts, real_samples.size(0)).cpu()
      generated_samples = torchcde.LinearInterpolation(
          generated_samples
      ).evaluate(ts)
      generated_samples = generated_samples[..., 1]
  
      # Plot histograms
      for prop in plot_locs:
          time = int(prop * (real_samples.size(1) - 1))
          real_samples_time = real_samples[:, time]
          generated_samples_time = generated_samples[:, time]
          _, bins, _ = plt.hist(
              real_samples_time.cpu().numpy(),
              bins=32,
              alpha=0.7,
              label="Real",
              color="dodgerblue",
              density=True,
          )
          bin_width = bins[1] - bins[0]
          num_bins = int(
              (
                  generated_samples_time.max()
                  - generated_samples_time.min()
              ).item()
              // bin_width
          )
          plt.hist(
              generated_samples_time.cpu().numpy(),
              bins=num_bins,
              alpha=0.7,
              label="Generated",
              color="crimson",
              density=True,
          )
          plt.legend()
          plt.xlabel("Value")
          plt.ylabel("Density")
          plt.title(f"Marginal distribution at time {time}.")
          plt.tight_layout()
          plt.show()
  
      real_samples = real_samples[:num_plot_samples]
      generated_samples = generated_samples[:num_plot_samples]
  
      # Plot samples
      real_first = True
      generated_first = True
      for real_sample_ in real_samples:
          kwargs = {"label": "Real"} if real_first else {}
          plt.plot(
              ts.cpu(),
              real_sample_.cpu(),
              color="dodgerblue",
              linewidth=0.5,
              alpha=0.7,
              **kwargs,
          )
          real_first = False
      for generated_sample_ in generated_samples:
          kwargs = {"label": "Generated"} if generated_first else {}
          plt.plot(
              ts.cpu(),
              generated_sample_.cpu(),
              color="crimson",
              linewidth=0.5,
              alpha=0.7,
              **kwargs,
          )
          generated_first = False
      plt.legend()
      plt.title(
          f"{num_plot_samples} samples from both real and generated distributions."
      )
      plt.tight_layout()
      plt.show()
  
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 2 for calling a global function with example usage: plot(ts, generator, dataloader, num_plot_samples, plot_locs) without return types