code: def weight_animate(i):
      ax.cla()
      ax.axis("off")
      ax.set_title("Routing: %d  " % i)
      dm = dist_list[i]
      nx.draw_networkx_nodes(
          g,
          pos,
          nodelist=range(in_nodes),
          node_color="r",
          node_size=100,
          ax=ax,
      )
      nx.draw_networkx_nodes(
          g,
          pos,
          nodelist=range(in_nodes, in_nodes + out_nodes),
          node_color="b",
          node_size=100,
          ax=ax,
      )
      for edge in g.edges():
          nx.draw_networkx_edges(
              g,
              pos,
              edgelist=[edge],
              width=dm[edge[0], edge[1] - in_nodes] * 1.5,
              ax=ax,
          )
  
file_path: dgl/tutorials/models/4_old_wines/2_capsule.py
imports: import os
  import dgl
  import matplotlib.pyplot
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import dgl.function
  import matplotlib.animation
  import seaborn
  import networkx
  from networkx.algorithms import bipartite
label_desc: Use image node_id 3 for calling a global function with example usage: weight_animate(i) without return types