code: def main(args):
      if args.cuda:
          torch.set_default_device("cuda")
  
      hidden_dim = args.hidden_dim
      obs_dim = args.obs_dim
      duration = args.duration
      batch_shape = (args.batch_size,)
  
      # Initialize parts.
      init_dist = random_mvn(
          batch_shape, hidden_dim, requires_grad=args.grad
      )
      trans_dist = random_mvn(
          batch_shape + (duration,), hidden_dim, requires_grad=args.grad
      )
      obs_dist = random_mvn(
          batch_shape + (1,), obs_dim, requires_grad=args.grad
      )
      trans_mat = 0.1 * torch.randn(
          batch_shape + (duration, hidden_dim, hidden_dim)
      )
      obs_mat = torch.randn(batch_shape + (1, hidden_dim, obs_dim))
  
      if args.grad:
          # Collect parameters.
          params = [
              init_dist.loc,
              init_dist.scale_tril,
              trans_dist.loc,
              trans_dist.scale_tril,
              obs_dist.loc,
              obs_dist.scale_tril,
              trans_mat.requires_grad_(),
              obs_mat.requires_grad_(),
          ]
  
      # Build a distribution.
      d = dist.GaussianHMM(
          init_dist,
          trans_mat,
          trans_dist,
          obs_mat,
          obs_dist,
          duration=duration,
      )
  
      for step in tqdm(range(args.num_steps)):
          if not args.grad:
              # Time forward only.
              d.sample()
              continue
  
          # Time forward + backward.
          x = d.rsample()
          grads = torch.autograd.grad(
              x.sum(), params, allow_unused=True, retain_graph=True
          )
          assert not all(g is None for g in grads)
          del x
  
imports: import argparse
  import torch
  from tqdm.auto import tqdm
  import pyro.distributions
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types