class_docstr: ## Configurations
  
  These are default configurations which can later be adjusted by passing a `dict`.
class_docstr_tok: ["#","#","Configurations","These","are","default","configurations","which","can","later","be","adjusted","by","passing","a","`","dict","`","."]
code: def step(self, batch: Any, batch_idx: BatchIndex):
      self.encoder.train(self.mode.is_train)
      self.decoder.train(self.mode.is_train)
  
      # Move `data` and `mask` to device and swap the sequence and batch dimensions.
      # `data` will have shape `[seq_len, batch_size, 5]` and
      # `mask` will have shape `[seq_len, batch_size]`.
      data = batch[0].to(self.device).transpose(0, 1)
      mask = batch[1].to(self.device).transpose(0, 1)
  
      # Increment step in training mode
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      # Encode the sequence of strokes
      with monit.section("encoder"):
          # Get $z$, $\mu$, and $\hat{\sigma}$
          z, mu, sigma_hat = self.encoder(data)
  
      # Decode the mixture of distributions and $\hat{q}$
      with monit.section("decoder"):
          # Concatenate $[(\Delta x, \Delta y, p_1, p_2, p_3); z]$
          z_stack = z.unsqueeze(0).expand(data.shape[0] - 1, -1, -1)
          inputs = torch.cat([data[:-1], z_stack], 2)
          # Get mixture of distributions and $\hat{q}$
          dist, q_logits, _ = self.decoder(inputs, z, None)
  
      # Compute the loss
      with monit.section("loss"):
          # $L_{KL}$
          kl_loss = self.kl_div_loss(sigma_hat, mu)
          # $L_R$
          reconstruction_loss = self.reconstruction_loss(
              mask, data[1:], dist, q_logits
          )
          # $Loss = L_R + w_{KL} L_{KL}$
          loss = reconstruction_loss + self.kl_div_loss_weight * kl_loss
  
          # Track losses
          tracker.add("loss.kl.", kl_loss)
          tracker.add("loss.reconstruction.", reconstruction_loss)
          tracker.add("loss.total.", loss)
  
      # Only if we are in training state
      if self.mode.is_train:
          # Run optimizer
          with monit.section("optimize"):
              # Set `grad` to zero
              self.optimizer.zero_grad()
              # Compute gradients
              loss.backward()
              # Log model parameters and gradients
              if batch_idx.is_last:
                  tracker.add(
                      encoder=self.encoder, decoder=self.decoder
                  )
              # Clip gradients
              nn.utils.clip_grad_norm_(
                  self.encoder.parameters(), self.grad_clip
              )
              nn.utils.clip_grad_norm_(
                  self.decoder.parameters(), self.grad_clip
              )
              # Optimize
              self.optimizer.step()
  
      tracker.save()
  
imports: import math
  from typing import Optional, Tuple, Any
  import numpy
  import torch
  import torch.nn
  from matplotlib import pyplot
  from torch import optim
  from torch.utils.data import Dataset, DataLoader
  import einops
  from labml import lab, experiment, tracker, monit
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.module import Module
  from labml_helpers.optimizer import OptimizerConfigs
  from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
label_desc: Use image node_id 2 for calling the Configs obj's underlying member method code with example usage: obj.step(batch, batch_idx) without return types