code: def main(
      train_dir="./dump/cont_ddpm/",
      epochs=100,
      lr=1e-4,
      batch_size=128,
      pause_every=1000,
      tau=1.0,
      logit_transform=True,
  ):
      """Train and sample once in a while.
  
      Args:
          train_dir: Path to a folder to dump things.
          epochs: Number of training epochs.
          lr: Learning rate for Adam.
          batch_size: Batch size for training.
          pause_every: Log and write figures once in this many iterations.
          tau: The temperature for sampling.
          logit_transform: Applies the typical logit transformation if True.
      """
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
  
      # Data.
      train_loader = make_loader(
          root=os.path.join(train_dir, "data"),
          train_batch_size=batch_size,
      )
  
      # Model + optimizer.
      denoiser = unet.Unet(
          input_size=(1, 28, 28),
          dim_mults=(
              1,
              2,
              4,
          ),
          attention_cls=unet.LinearTimeSelfAttention,
      )
      forward = ScoreMatchingSDE(denoiser=denoiser).to(device)
      reverse = ReverseDiffeqWrapper(forward)
      optimizer = optim.Adam(params=forward.parameters(), lr=lr)
  
      def plot(imgs, path):
          assert not torch.any(
              torch.isnan(imgs)
          ), "Found nans in images"
          os.makedirs(os.path.dirname(path), exist_ok=True)
          imgs = (
              postprocess(imgs, logit_transform=logit_transform)
              .detach()
              .cpu()
          )
          tv.utils.save_image(imgs, path)
  
      global_step = 0
      for epoch in range(epochs):
          for x, _ in tqdm.tqdm(train_loader):
              forward.train()
              forward.zero_grad()
              x = preprocess(
                  x.to(device), logit_transform=logit_transform
              )
              loss = forward(x).mean(dim=0)
              loss.backward()
              optimizer.step()
              global_step += 1
  
              if global_step % pause_every == 0:
                  logging.warning(
                      f"global_step: {global_step:06d}, loss: {loss:.4f}"
                  )
  
                  img_path = os.path.join(
                      train_dir,
                      "ode_samples",
                      f"global_step_{global_step:07d}.png",
                  )
                  ode_samples = reverse.ode_sample_final(tau=tau)
                  plot(ode_samples, img_path)
  
                  img_path = os.path.join(
                      train_dir,
                      "sde_samples",
                      f"global_step_{global_step:07d}.png",
                  )
                  sde_samples = reverse.sde_sample_final(tau=tau)
                  plot(sde_samples, img_path)
  
imports: import abc
  import logging
  import math
  import os
  import fire
  import torch
  import torchdiffeq
  import torchvision
  import tqdm
  from torch import nn, optim
  from torch.utils import data
  import torchsde
  from .None import unet
label_desc: Use image node_id 5 for calling a global function with example usage: main(train_dir, epochs, lr, batch_size, pause_every, tau, logit_transform) without return types