From 34d16cf6fb97b21376e0503964f7f0a8082709d4 Mon Sep 17 00:00:00 2001
From: Aarav Pandya <ap7641@nyu.edu>
Date: Fri, 12 Jan 2024 20:39:04 -0500
Subject: [PATCH 1/2] Init

---
 examples/06_ppo_with_sb3_ma_menv_control.py |  193 +
 examples/hr_rl.ipynb                        | 6640 +++++++++++++++++++
 examples/temp.ipynb                         |  384 ++
 nocturne/envs/nocturne_gymnasium.py         |  136 +
 4 files changed, 7353 insertions(+)
 create mode 100644 examples/06_ppo_with_sb3_ma_menv_control.py
 create mode 100644 examples/hr_rl.ipynb
 create mode 100644 examples/temp.ipynb
 create mode 100644 nocturne/envs/nocturne_gymnasium.py

diff --git a/examples/06_ppo_with_sb3_ma_menv_control.py b/examples/06_ppo_with_sb3_ma_menv_control.py
new file mode 100644
index 00000000..2f82e87d
--- /dev/null
+++ b/examples/06_ppo_with_sb3_ma_menv_control.py
@@ -0,0 +1,193 @@
+"""Train HR-PPO agent."""
+import logging
+from contextlib import nullcontext
+from datetime import datetime
+
+import numpy as np
+import torch
+from box import Box
+from stable_baselines3.common.policies import ActorCriticPolicy
+
+import wandb
+
+from typing import Callable
+
+# Import networks
+from networks.mlp_late_fusion import LateFusionMLP, LateFusionMLPPolicy
+# Permutation equivariant network
+from networks.perm_eq_late_fusion import LateFusionNet, LateFusionPolicy 
+
+# Multi-agent as vectorized environment
+from nocturne.envs.vec_env_ma import MultiAgentAsVecEnv
+from utils.config import load_config
+from utils.random_utils import init_seed
+from utils.render import make_video
+
+# Custom callback
+from utils.sb3.callbacks import CustomMultiAgentCallback
+
+# Custom PPO class that supports multi-agent control
+from utils.sb3.reg_ppo import RegularizedPPO
+from utils.string_utils import datetime_to_str
+
+logging.basicConfig(level=logging.INFO)
+
+
+def linear_schedule(initial_value: float) -> Callable[[float], float]:
+    """
+    Linear learning rate schedule.
+
+    :param initial_value: Initial learning rate.
+    :return: schedule that computes
+      current learning rate depending on remaining progress
+    """
+    def func(progress_remaining: float) -> float:
+        """
+        Progress will decrease from 1 (beginning) to 0.
+
+        :param progress_remaining:
+        :return: current learning rate
+        """
+        return progress_remaining * initial_value
+
+    return func
+
+def train(env_config, exp_config, video_config, model_config):  # pylint: disable=redefined-outer-name
+    """Train RL agent using PPO."""
+    # Ensure reproducability
+    init_seed(env_config, exp_config, exp_config.seed)
+
+    # Make environment
+    from nocturne.envs.nocturne_gymnasium import NocturneGymnasium
+    from stable_baselines3.common.vec_env import SubprocVecEnv
+    from nocturne.envs.base_env import BaseEnv
+    def make_env(env_config):
+        return NocturneGymnasium(BaseEnv(config=env_config)) 
+
+    env = SubprocVecEnv([lambda: make_env(env_config) for _ in range(4)])
+
+    # Set up run
+    datetime_ = datetime_to_str(dt=datetime.now())
+    run_id = f"{datetime_}" if exp_config.track_wandb else None
+
+    # Add scene to config
+    # exp_config.scene = env.filename
+    exp_config.track_wandb = False
+
+    with wandb.init(
+        project=exp_config.project,
+        name=run_id,
+        group=exp_config.group,
+        config={**exp_config, **env_config},
+        id=run_id,
+        **exp_config.wandb,
+    ) if exp_config.track_wandb else nullcontext() as run:
+        # Set device
+        exp_config.ppo.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # logging.info(f"Created env. Max # agents = {env_config.max_num_vehicles}.")
+        # logging.info(f"Learning in {len(env.env.files)} scene(s): {env.env.files} | using {exp_config.ppo.device}")
+        # logging.info(f"--- obs_space: {env.observation_space.shape[0]} ---")
+        # logging.info(f"Action_space\n: {env.env.idx_to_actions}")
+        
+        # if exp_config.reg_weight > 0.0:
+        #     logging.info(f"Regularization weight: {exp_config.reg_weight} with policy: {exp_config.human_policy_path}")
+
+        # # Initialize custom callback
+        custom_callback = CustomMultiAgentCallback(
+            env_config=env_config,
+            exp_config=exp_config,
+            video_config=video_config,
+            wandb_run=run if run_id is not None else None,
+        )
+
+        # Make scene init video to check expert actions
+        # if exp_config.track_wandb:
+        #     for model in exp_config.wandb_init_videos:
+        #         make_video(
+        #             env_config=env_config,
+        #             exp_config=exp_config,
+        #             video_config=video_config,
+        #             filenames=[env.filename],
+        #             model=model,
+        #             n_steps=None,
+        #         )
+        exp_config.track_wandb = False
+        
+        human_policy = None
+        # Load human reference policy if regularization is used
+        if exp_config.reg_weight > 0.0:
+            saved_variables = torch.load(exp_config.human_policy_path, map_location=exp_config.ppo.device)
+            human_policy = ActorCriticPolicy(**saved_variables["data"])
+            human_policy.load_state_dict(saved_variables["state_dict"])
+            human_policy.to(exp_config.ppo.device)
+
+        # Set up PPO
+        model = RegularizedPPO(
+            learning_rate=linear_schedule(1e-4),
+            reg_policy=human_policy,
+            reg_weight=exp_config.reg_weight,  # Regularization weight; lambda
+            env=env,
+            n_steps=exp_config.ppo.n_steps,
+            policy=LateFusionPolicy,
+            ent_coef=exp_config.ppo.ent_coef,
+            vf_coef=exp_config.ppo.vf_coef,
+            seed=exp_config.seed,  # Seed for the pseudo random generators
+            verbose=exp_config.verbose,
+            tensorboard_log=f"runs/{run_id}" if run_id is not None else None,
+            device=exp_config.ppo.device,
+            env_config=env_config,
+            mlp_class=LateFusionNet,
+            mlp_config=model_config,
+        )
+
+        # Log number of trainable parameters
+        policy_params = filter(lambda p: p.requires_grad, model.policy.parameters())
+        params = sum(np.prod(p.size()) for p in policy_params)
+        exp_config.n_policy_params = params
+        logging.info(f"Policy | trainable params: {params:,} \n")
+
+        # Architecture
+        logging.info(f"Policy | arch: \n {model.policy}")
+
+
+        return env, model
+        # Learn
+        # model.learn(
+        #     **exp_config.learn,
+        #     callback=custom_callback,
+        # )
+
+
+if __name__ == "__main__":
+    env_config = load_config("env_config")
+    exp_config = load_config("exp_config")
+    video_config = load_config("video_config")
+
+    env_config.num_files = 10
+
+    # Define model architecture
+    model_config = None
+    # model_config = Box(
+    #     {
+    #         "arch_ego_state": [8],
+    #         "arch_road_objects": [64],
+    #         "arch_road_graph": [128, 64],
+    #         "arch_shared_net": [128],
+    #         "act_func": "tanh",
+    #         "dropout": 0.0,
+    #         "last_layer_dim_pi": 64,
+    #         "last_layer_dim_vf": 64,
+    #     }
+    # )
+
+    # Train
+    env, model = train(
+        env_config=env_config,
+        exp_config=exp_config,
+        video_config=video_config,
+        model_config=model_config,
+    )
+
+
+    model.learn(1)
\ No newline at end of file
diff --git a/examples/hr_rl.ipynb b/examples/hr_rl.ipynb
new file mode 100644
index 00000000..94ed4698
--- /dev/null
+++ b/examples/hr_rl.ipynb
@@ -0,0 +1,6640 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"Train HR-PPO agent.\"\"\"\n",
+    "import logging\n",
+    "from contextlib import nullcontext\n",
+    "from datetime import datetime\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from box import Box\n",
+    "from stable_baselines3.common.policies import ActorCriticPolicy\n",
+    "\n",
+    "import wandb\n",
+    "\n",
+    "from typing import Callable\n",
+    "\n",
+    "# Import networks\n",
+    "from networks.mlp_late_fusion import LateFusionMLP, LateFusionMLPPolicy\n",
+    "# Permutation equivariant network\n",
+    "from networks.perm_eq_late_fusion import LateFusionNet, LateFusionPolicy \n",
+    "\n",
+    "# Multi-agent as vectorized environment\n",
+    "from nocturne.envs.vec_env_ma import MultiAgentAsVecEnv\n",
+    "from utils.config import load_config_nb\n",
+    "from utils.random_utils import init_seed\n",
+    "from utils.render import make_video\n",
+    "\n",
+    "# Custom callback\n",
+    "from utils.sb3.callbacks import CustomMultiAgentCallback\n",
+    "\n",
+    "# Custom PPO class that supports multi-agent control\n",
+    "from utils.sb3.reg_ppo import RegularizedPPO\n",
+    "from utils.string_utils import datetime_to_str\n",
+    "\n",
+    "logging.basicConfig(level=logging.INFO)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def linear_schedule(initial_value: float) -> Callable[[float], float]:\n",
+    "    \"\"\"\n",
+    "    Linear learning rate schedule.\n",
+    "\n",
+    "    :param initial_value: Initial learning rate.\n",
+    "    :return: schedule that computes\n",
+    "      current learning rate depending on remaining progress\n",
+    "    \"\"\"\n",
+    "    def func(progress_remaining: float) -> float:\n",
+    "        \"\"\"\n",
+    "        Progress will decrease from 1 (beginning) to 0.\n",
+    "\n",
+    "        :param progress_remaining:\n",
+    "        :return: current learning rate\n",
+    "        \"\"\"\n",
+    "        return progress_remaining * initial_value\n",
+    "\n",
+    "    return func\n",
+    "\n",
+    "def train(env_config, exp_config, video_config, model_config):  # pylint: disable=redefined-outer-name\n",
+    "    \"\"\"Train RL agent using PPO.\"\"\"\n",
+    "    # Ensure reproducability\n",
+    "    init_seed(env_config, exp_config, exp_config.seed)\n",
+    "\n",
+    "    # Make environment\n",
+    "    from nocturne.envs.nocturne_gymnasium import NocturneGymnasium\n",
+    "    from stable_baselines3.common.vec_env import SubprocVecEnv\n",
+    "    from nocturne.envs.base_env import BaseEnv\n",
+    "    def make_env(env_config):\n",
+    "        return NocturneGymnasium(BaseEnv(config=env_config)) \n",
+    "\n",
+    "    env = SubprocVecEnv([lambda: make_env(env_config) for _ in range(4)])\n",
+    "\n",
+    "    # Set up run\n",
+    "    datetime_ = datetime_to_str(dt=datetime.now())\n",
+    "    run_id = f\"{datetime_}\" if exp_config.track_wandb else None\n",
+    "\n",
+    "    # Add scene to config\n",
+    "    # exp_config.scene = env.filename\n",
+    "    exp_config.track_wandb = False\n",
+    "\n",
+    "    with wandb.init(\n",
+    "        project=exp_config.project,\n",
+    "        name=run_id,\n",
+    "        group=exp_config.group,\n",
+    "        config={**exp_config, **env_config},\n",
+    "        id=run_id,\n",
+    "        **exp_config.wandb,\n",
+    "    ) if exp_config.track_wandb else nullcontext() as run:\n",
+    "        # Set device\n",
+    "        exp_config.ppo.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "        # logging.info(f\"Created env. Max # agents = {env_config.max_num_vehicles}.\")\n",
+    "        # logging.info(f\"Learning in {len(env.env.files)} scene(s): {env.env.files} | using {exp_config.ppo.device}\")\n",
+    "        # logging.info(f\"--- obs_space: {env.observation_space.shape[0]} ---\")\n",
+    "        # logging.info(f\"Action_space\\n: {env.env.idx_to_actions}\")\n",
+    "        \n",
+    "        # if exp_config.reg_weight > 0.0:\n",
+    "        #     logging.info(f\"Regularization weight: {exp_config.reg_weight} with policy: {exp_config.human_policy_path}\")\n",
+    "\n",
+    "        # # Initialize custom callback\n",
+    "        custom_callback = CustomMultiAgentCallback(\n",
+    "            env_config=env_config,\n",
+    "            exp_config=exp_config,\n",
+    "            video_config=video_config,\n",
+    "            wandb_run=run if run_id is not None else None,\n",
+    "        )\n",
+    "\n",
+    "        # Make scene init video to check expert actions\n",
+    "        # if exp_config.track_wandb:\n",
+    "        #     for model in exp_config.wandb_init_videos:\n",
+    "        #         make_video(\n",
+    "        #             env_config=env_config,\n",
+    "        #             exp_config=exp_config,\n",
+    "        #             video_config=video_config,\n",
+    "        #             filenames=[env.filename],\n",
+    "        #             model=model,\n",
+    "        #             n_steps=None,\n",
+    "        #         )\n",
+    "        exp_config.track_wandb = False\n",
+    "        \n",
+    "        human_policy = None\n",
+    "        # Load human reference policy if regularization is used\n",
+    "        if exp_config.reg_weight > 0.0:\n",
+    "            saved_variables = torch.load(exp_config.human_policy_path, map_location=exp_config.ppo.device)\n",
+    "            human_policy = ActorCriticPolicy(**saved_variables[\"data\"])\n",
+    "            human_policy.load_state_dict(saved_variables[\"state_dict\"])\n",
+    "            human_policy.to(exp_config.ppo.device)\n",
+    "\n",
+    "        # Set up PPO\n",
+    "        model = RegularizedPPO(\n",
+    "            learning_rate=linear_schedule(1e-4),\n",
+    "            reg_policy=human_policy,\n",
+    "            reg_weight=exp_config.reg_weight,  # Regularization weight; lambda\n",
+    "            env=env,\n",
+    "            n_steps=exp_config.ppo.n_steps,\n",
+    "            policy=LateFusionPolicy,\n",
+    "            ent_coef=exp_config.ppo.ent_coef,\n",
+    "            vf_coef=exp_config.ppo.vf_coef,\n",
+    "            seed=exp_config.seed,  # Seed for the pseudo random generators\n",
+    "            verbose=exp_config.verbose,\n",
+    "            tensorboard_log=f\"runs/{run_id}\" if run_id is not None else None,\n",
+    "            device=exp_config.ppo.device,\n",
+    "            env_config=env_config,\n",
+    "            mlp_class=LateFusionNet,\n",
+    "            mlp_config=model_config,\n",
+    "        )\n",
+    "\n",
+    "        # Log number of trainable parameters\n",
+    "        policy_params = filter(lambda p: p.requires_grad, model.policy.parameters())\n",
+    "        params = sum(np.prod(p.size()) for p in policy_params)\n",
+    "        exp_config.n_policy_params = params\n",
+    "        logging.info(f\"Policy | trainable params: {params:,} \\n\")\n",
+    "\n",
+    "        # Architecture\n",
+    "        logging.info(f\"Policy | arch: \\n {model.policy}\")\n",
+    "\n",
+    "\n",
+    "        return env, model\n",
+    "        # Learn\n",
+    "        # model.learn(\n",
+    "        #     **exp_config.learn,\n",
+    "        #     callback=custom_callback,\n",
+    "        # )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:root:Policy | trainable params: 148,566 \n",
+      "\n",
+      "INFO:root:Policy | arch: \n",
+      " LateFusionPolicy(\n",
+      "  (features_extractor): FlattenExtractor(\n",
+      "    (flatten): Flatten(start_dim=1, end_dim=-1)\n",
+      "  )\n",
+      "  (pi_features_extractor): FlattenExtractor(\n",
+      "    (flatten): Flatten(start_dim=1, end_dim=-1)\n",
+      "  )\n",
+      "  (vf_features_extractor): FlattenExtractor(\n",
+      "    (flatten): Flatten(start_dim=1, end_dim=-1)\n",
+      "  )\n",
+      "  (mlp_extractor): LateFusionNet(\n",
+      "    (act_func): Tanh()\n",
+      "    (actor_ego_state_net): Sequential(\n",
+      "      (0): Linear(in_features=10, out_features=10, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((10,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "    )\n",
+      "    (actor_ro_net): Sequential(\n",
+      "      (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "      (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "      (5): Dropout(p=0.0, inplace=False)\n",
+      "      (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "      (7): Tanh()\n",
+      "    )\n",
+      "    (actor_rg_net): Sequential(\n",
+      "      (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "      (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "      (5): Dropout(p=0.0, inplace=False)\n",
+      "      (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "      (7): Tanh()\n",
+      "    )\n",
+      "    (actor_ss_net): Sequential(\n",
+      "      (0): Linear(in_features=3, out_features=3, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((3,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "    )\n",
+      "    (actor_out_net): Sequential(\n",
+      "      (0): Linear(in_features=77, out_features=256, bias=True)\n",
+      "      (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
+      "      (2): Tanh()\n",
+      "      (3): Dropout(p=0.0, inplace=False)\n",
+      "      (4): Linear(in_features=256, out_features=128, bias=True)\n",
+      "      (5): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
+      "      (6): Tanh()\n",
+      "      (7): Dropout(p=0.0, inplace=False)\n",
+      "      (8): Linear(in_features=128, out_features=64, bias=True)\n",
+      "      (9): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "      (10): Tanh()\n",
+      "      (11): Dropout(p=0.0, inplace=False)\n",
+      "      (12): Linear(in_features=64, out_features=64, bias=True)\n",
+      "      (13): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    )\n",
+      "    (val_ego_state_net): Sequential(\n",
+      "      (0): Linear(in_features=10, out_features=10, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((10,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "    )\n",
+      "    (val_ro_net): Sequential(\n",
+      "      (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "      (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "      (5): Dropout(p=0.0, inplace=False)\n",
+      "      (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "      (7): Tanh()\n",
+      "    )\n",
+      "    (val_rg_net): Sequential(\n",
+      "      (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "      (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "      (5): Dropout(p=0.0, inplace=False)\n",
+      "      (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "      (7): Tanh()\n",
+      "    )\n",
+      "    (val_ss_net): Sequential(\n",
+      "      (0): Linear(in_features=3, out_features=3, bias=True)\n",
+      "      (1): Dropout(p=0.0, inplace=False)\n",
+      "      (2): LayerNorm((3,), eps=1e-05, elementwise_affine=True)\n",
+      "      (3): Tanh()\n",
+      "    )\n",
+      "    (val_out_net): Sequential(\n",
+      "      (0): Linear(in_features=77, out_features=256, bias=True)\n",
+      "      (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
+      "      (2): Tanh()\n",
+      "      (3): Dropout(p=0.0, inplace=False)\n",
+      "      (4): Linear(in_features=256, out_features=128, bias=True)\n",
+      "      (5): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
+      "      (6): Tanh()\n",
+      "      (7): Dropout(p=0.0, inplace=False)\n",
+      "      (8): Linear(in_features=128, out_features=64, bias=True)\n",
+      "      (9): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "      (10): Tanh()\n",
+      "      (11): Dropout(p=0.0, inplace=False)\n",
+      "      (12): Linear(in_features=64, out_features=64, bias=True)\n",
+      "      (13): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    )\n",
+      "  )\n",
+      "  (action_net): Linear(in_features=64, out_features=45, bias=True)\n",
+      "  (value_net): Linear(in_features=64, out_features=1, bias=True)\n",
+      ")\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "LateFusionNet(\n",
+      "  (act_func): Tanh()\n",
+      "  (actor_ego_state_net): Sequential(\n",
+      "    (0): Linear(in_features=10, out_features=10, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((10,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "  )\n",
+      "  (actor_ro_net): Sequential(\n",
+      "    (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "    (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "    (5): Dropout(p=0.0, inplace=False)\n",
+      "    (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "    (7): Tanh()\n",
+      "  )\n",
+      "  (actor_rg_net): Sequential(\n",
+      "    (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "    (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "    (5): Dropout(p=0.0, inplace=False)\n",
+      "    (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "    (7): Tanh()\n",
+      "  )\n",
+      "  (actor_ss_net): Sequential(\n",
+      "    (0): Linear(in_features=3, out_features=3, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((3,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "  )\n",
+      "  (actor_out_net): Sequential(\n",
+      "    (0): Linear(in_features=77, out_features=256, bias=True)\n",
+      "    (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
+      "    (2): Tanh()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=256, out_features=128, bias=True)\n",
+      "    (5): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
+      "    (6): Tanh()\n",
+      "    (7): Dropout(p=0.0, inplace=False)\n",
+      "    (8): Linear(in_features=128, out_features=64, bias=True)\n",
+      "    (9): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    (10): Tanh()\n",
+      "    (11): Dropout(p=0.0, inplace=False)\n",
+      "    (12): Linear(in_features=64, out_features=64, bias=True)\n",
+      "    (13): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "  )\n",
+      "  (val_ego_state_net): Sequential(\n",
+      "    (0): Linear(in_features=10, out_features=10, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((10,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "  )\n",
+      "  (val_ro_net): Sequential(\n",
+      "    (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "    (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "    (5): Dropout(p=0.0, inplace=False)\n",
+      "    (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "    (7): Tanh()\n",
+      "  )\n",
+      "  (val_rg_net): Sequential(\n",
+      "    (0): Linear(in_features=13, out_features=64, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "    (4): Linear(in_features=64, out_features=32, bias=True)\n",
+      "    (5): Dropout(p=0.0, inplace=False)\n",
+      "    (6): LayerNorm((32,), eps=1e-05, elementwise_affine=True)\n",
+      "    (7): Tanh()\n",
+      "  )\n",
+      "  (val_ss_net): Sequential(\n",
+      "    (0): Linear(in_features=3, out_features=3, bias=True)\n",
+      "    (1): Dropout(p=0.0, inplace=False)\n",
+      "    (2): LayerNorm((3,), eps=1e-05, elementwise_affine=True)\n",
+      "    (3): Tanh()\n",
+      "  )\n",
+      "  (val_out_net): Sequential(\n",
+      "    (0): Linear(in_features=77, out_features=256, bias=True)\n",
+      "    (1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
+      "    (2): Tanh()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=256, out_features=128, bias=True)\n",
+      "    (5): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
+      "    (6): Tanh()\n",
+      "    (7): Dropout(p=0.0, inplace=False)\n",
+      "    (8): Linear(in_features=128, out_features=64, bias=True)\n",
+      "    (9): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "    (10): Tanh()\n",
+      "    (11): Dropout(p=0.0, inplace=False)\n",
+      "    (12): Linear(in_features=64, out_features=64, bias=True)\n",
+      "    (13): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
+      "  )\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "env_config = load_config_nb(\"env_config\")\n",
+    "exp_config = load_config_nb(\"exp_config\")\n",
+    "video_config = load_config_nb(\"video_config\")\n",
+    "\n",
+    "env_config.num_files = 10\n",
+    "\n",
+    "# Define model architecture\n",
+    "model_config = None\n",
+    "# model_config = Box(\n",
+    "#     {\n",
+    "#         \"arch_ego_state\": [8],\n",
+    "#         \"arch_road_objects\": [64],\n",
+    "#         \"arch_road_graph\": [128, 64],\n",
+    "#         \"arch_shared_net\": [128],\n",
+    "#         \"act_func\": \"tanh\",\n",
+    "#         \"dropout\": 0.0,\n",
+    "#         \"last_layer_dim_pi\": 64,\n",
+    "#         \"last_layer_dim_vf\": 64,\n",
+    "#     }\n",
+    "# )\n",
+    "\n",
+    "# Train\n",
+    "env, model = train(\n",
+    "    env_config=env_config,\n",
+    "    exp_config=exp_config,\n",
+    "    video_config=video_config,\n",
+    "    model_config=model_config,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([{37: array([0.28523394, 0.49977985, 0.10757028, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 0: array([0.2631481 , 0.50686294, 0.01567911, ..., 0.01      , 0.68215334,\n",
+       "              0.01584669]), 51: array([0.33037499, 0.583     , 0.02133806, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 31: array([0.29642355, 0.51919246, 0.05890008, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 34: array([0.29417717, 0.51871204, 0.11663519, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 2: array([0.29118758, 0.50871187, 0.02653475, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 1: array([0.28495339, 0.51357961, 0.02791023, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 32: array([0.29227659, 0.5093497 , 0.11704891, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 30: array([0.28525367, 0.50776368, 0.125852  , ..., 0.        , 0.        ,\n",
+       "              0.        ]), 33: array([0.3665545 , 0.60222054, 0.10613482, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 41: array([0.27929321, 0.49081314, 0.11201638, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 4: array([0.29417551, 0.52149814, 0.11362467, ..., 0.        , 0.        ,\n",
+       "              0.        ])}                                                                            ,\n",
+       "       {37: array([0.28523394, 0.49977985, 0.10757028, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 0: array([0.2631481 , 0.50686294, 0.01567911, ..., 0.01      , 0.68215334,\n",
+       "              0.01584669]), 51: array([0.33037499, 0.583     , 0.02133806, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 31: array([0.29642355, 0.51919246, 0.05890008, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 34: array([0.29417717, 0.51871204, 0.11663519, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 2: array([0.29118758, 0.50871187, 0.02653475, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 1: array([0.28495339, 0.51357961, 0.02791023, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 32: array([0.29227659, 0.5093497 , 0.11704891, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 30: array([0.28525367, 0.50776368, 0.125852  , ..., 0.        , 0.        ,\n",
+       "              0.        ]), 33: array([0.3665545 , 0.60222054, 0.10613482, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 41: array([0.27929321, 0.49081314, 0.11201638, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 4: array([0.29417551, 0.52149814, 0.11362467, ..., 0.        , 0.        ,\n",
+       "              0.        ])}                                                                            ,\n",
+       "       {37: array([0.28523394, 0.49977985, 0.10757028, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 0: array([0.2631481 , 0.50686294, 0.01567911, ..., 0.01      , 0.68215334,\n",
+       "              0.01584669]), 51: array([0.33037499, 0.583     , 0.02133806, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 31: array([0.29642355, 0.51919246, 0.05890008, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 34: array([0.29417717, 0.51871204, 0.11663519, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 2: array([0.29118758, 0.50871187, 0.02653475, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 1: array([0.28495339, 0.51357961, 0.02791023, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 32: array([0.29227659, 0.5093497 , 0.11704891, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 30: array([0.28525367, 0.50776368, 0.125852  , ..., 0.        , 0.        ,\n",
+       "              0.        ]), 33: array([0.3665545 , 0.60222054, 0.10613482, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 41: array([0.27929321, 0.49081314, 0.11201638, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 4: array([0.29417551, 0.52149814, 0.11362467, ..., 0.        , 0.        ,\n",
+       "              0.        ])}                                                                            ,\n",
+       "       {37: array([0.28523394, 0.49977985, 0.10757028, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 0: array([0.2631481 , 0.50686294, 0.01567911, ..., 0.01      , 0.68215334,\n",
+       "              0.01584669]), 51: array([0.33037499, 0.583     , 0.02133806, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 31: array([0.29642355, 0.51919246, 0.05890008, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 34: array([0.29417717, 0.51871204, 0.11663519, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 2: array([0.29118758, 0.50871187, 0.02653475, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 1: array([0.28495339, 0.51357961, 0.02791023, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 32: array([0.29227659, 0.5093497 , 0.11704891, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 30: array([0.28525367, 0.50776368, 0.125852  , ..., 0.        , 0.        ,\n",
+       "              0.        ]), 33: array([0.3665545 , 0.60222054, 0.10613482, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 41: array([0.27929321, 0.49081314, 0.11201638, ..., 0.        , 0.        ,\n",
+       "              0.        ]), 4: array([0.29417551, 0.52149814, 0.11362467, ..., 0.        , 0.        ,\n",
+       "              0.        ])}                                                                            ],\n",
+       "      dtype=object)"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "env.reset()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "TypeError",
+     "evalue": "can't convert np.ndarray of type numpy.object_. The only supported types are: float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[6], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/nocturne_lab/.venv/lib/python3.11/site-packages/stable_baselines3/ppo/ppo.py:315\u001b[0m, in \u001b[0;36mPPO.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m    306\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mlearn\u001b[39m(\n\u001b[1;32m    307\u001b[0m     \u001b[38;5;28mself\u001b[39m: SelfPPO,\n\u001b[1;32m    308\u001b[0m     total_timesteps: \u001b[38;5;28mint\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    313\u001b[0m     progress_bar: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    314\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m SelfPPO:\n\u001b[0;32m--> 315\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    316\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    317\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    318\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_interval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    319\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_log_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    320\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    321\u001b[0m \u001b[43m        \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    322\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m~/nocturne_lab/.venv/lib/python3.11/site-packages/stable_baselines3/common/on_policy_algorithm.py:277\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.learn\u001b[0;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[1;32m    274\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39menv \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    276\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m<\u001b[39m total_timesteps:\n\u001b[0;32m--> 277\u001b[0m     continue_training \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect_rollouts\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrollout_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_rollout_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_steps\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    279\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m continue_training:\n\u001b[1;32m    280\u001b[0m         \u001b[38;5;28;01mbreak\u001b[39;00m\n",
+      "File \u001b[0;32m~/nocturne_lab/utils/sb3/ma_ppo.py:56\u001b[0m, in \u001b[0;36mMultiAgentPPO.collect_rollouts\u001b[0;34m(self, env, callback, rollout_buffer, n_rollout_steps)\u001b[0m\n\u001b[1;32m     52\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpolicy\u001b[38;5;241m.\u001b[39mreset_noise(env\u001b[38;5;241m.\u001b[39mnum_envs)\n\u001b[1;32m     54\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m     55\u001b[0m     \u001b[38;5;66;03m# Convert to pytorch tensor or to TensorDict\u001b[39;00m\n\u001b[0;32m---> 56\u001b[0m     obs_tensor \u001b[38;5;241m=\u001b[39m \u001b[43mobs_as_tensor\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_last_obs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     58\u001b[0m     \u001b[38;5;66;03m# EDIT_1: Mask out invalid observations (NaN dimensions and/or dead agents)\u001b[39;00m\n\u001b[1;32m     59\u001b[0m     \u001b[38;5;66;03m# Create dummy actions, values and log_probs (NaN)\u001b[39;00m\n\u001b[1;32m     60\u001b[0m     actions \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mfull(fill_value\u001b[38;5;241m=\u001b[39mnp\u001b[38;5;241m.\u001b[39mnan, size\u001b[38;5;241m=\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn_envs,))\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice)\n",
+      "File \u001b[0;32m~/nocturne_lab/.venv/lib/python3.11/site-packages/stable_baselines3/common/utils.py:483\u001b[0m, in \u001b[0;36mobs_as_tensor\u001b[0;34m(obs, device)\u001b[0m\n\u001b[1;32m    475\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    476\u001b[0m \u001b[38;5;124;03mMoves the observation to the given device.\u001b[39;00m\n\u001b[1;32m    477\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    480\u001b[0m \u001b[38;5;124;03m:return: PyTorch tensor of the observation on a desired device.\u001b[39;00m\n\u001b[1;32m    481\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    482\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(obs, np\u001b[38;5;241m.\u001b[39mndarray):\n\u001b[0;32m--> 483\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mth\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mas_tensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    484\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(obs, \u001b[38;5;28mdict\u001b[39m):\n\u001b[1;32m    485\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m {key: th\u001b[38;5;241m.\u001b[39mas_tensor(_obs, device\u001b[38;5;241m=\u001b[39mdevice) \u001b[38;5;28;01mfor\u001b[39;00m (key, _obs) \u001b[38;5;129;01min\u001b[39;00m obs\u001b[38;5;241m.\u001b[39mitems()}\n",
+      "\u001b[0;31mTypeError\u001b[0m: can't convert np.ndarray of type numpy.object_. The only supported types are: float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool."
+     ]
+    }
+   ],
+   "source": [
+    "model.learn(1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "y = env.observation_space.sample()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Discrete(25)"
+      ]
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "env.action_space"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "m = model.policy.eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(tensor([8], device='cuda:0'),\n",
+       " tensor([[-0.5645]], device='cuda:0', grad_fn=<AddmmBackward0>),\n",
+       " tensor([-3.6230], device='cuda:0', grad_fn=<SqueezeBackward1>))"
+      ]
+     },
+     "execution_count": 50,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "m(torch.Tensor(y).unsqueeze(0).cuda())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torchviz import make_dot"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(tensor([[-0.6441, -1.3309,  0.5179, -0.4558,  0.6692,  0.3656,  1.1654,  1.2158,\n",
+       "           0.9431, -0.3549, -1.8127, -0.9146,  0.6600,  1.6707, -0.7357,  1.0185,\n",
+       "           1.4673,  0.1360,  1.5719, -0.4908, -0.5020, -0.1510, -2.7177, -0.9720,\n",
+       "          -1.5061,  0.1968,  0.1585,  0.3488, -0.5726,  0.5933, -0.0064, -1.4018,\n",
+       "           0.3319, -2.0170,  0.8157, -1.6580,  0.5165,  0.6804, -1.1788, -0.5321,\n",
+       "          -0.3591,  0.4718,  0.4771,  0.6652,  0.6524,  0.5661, -1.1508,  0.2073,\n",
+       "          -0.4792,  1.4354,  1.1228,  1.2609,  0.6219, -1.6648,  0.5761,  0.1294,\n",
+       "           0.4942, -1.1621,  1.7667, -1.1287,  0.1603,  0.8658,  0.2625, -0.8793]],\n",
+       "        device='cuda:0', grad_fn=<NativeLayerNormBackward0>),\n",
+       " tensor([[-0.3148,  1.7813,  3.3665,  0.3266, -0.4068,  0.3553, -1.4078, -1.1260,\n",
+       "           0.0797,  0.4915,  0.1167, -0.4714, -0.7219,  1.0971, -0.2553, -0.5615,\n",
+       "           0.3163,  1.6734, -0.1957,  0.0232, -0.9513,  0.4992, -0.1210, -0.9332,\n",
+       "           0.1590, -1.3791,  0.4740, -0.7048, -0.2949, -0.6597,  0.9464,  0.3830,\n",
+       "          -2.3730, -0.4302, -0.7398, -0.3818,  0.6071,  1.6034,  1.2792, -1.6859,\n",
+       "           0.6433,  0.4296, -2.0910,  0.2231,  0.6841,  0.2559, -0.5256, -0.4875,\n",
+       "           2.0396, -0.4128,  0.8378, -0.6700, -0.4007,  0.2891, -1.4553, -1.0730,\n",
+       "          -0.8699,  1.3477,  0.5756,  1.1768, -0.5317,  0.2053,  0.5588, -0.2122]],\n",
+       "        device='cuda:0', grad_fn=<NativeLayerNormBackward0>))"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.policy.mlp_extractor.cuda()(torch.Tensor(y).unsqueeze(0).cuda())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# make_dot(model.policy(torch.Tensor(y).unsqueeze(0).cuda()), params=dict(model.policy.named_parameters()))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torchview import draw_graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[-0.2501,  0.3726,  0.1859,  ..., -0.5912,  0.7823, -1.0088]])"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "torch.Tensor(y).unsqueeze(0).cpu()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/svg+xml": [
+       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
+       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
+       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
+       "<!-- Generated by graphviz version 2.43.0 (0)\n",
+       " -->\n",
+       "<!-- Title: model Pages: 1 -->\n",
+       "<svg width=\"1852pt\" height=\"2370pt\"\n",
+       " viewBox=\"0.00 0.00 1851.50 2370.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
+       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 2366)\">\n",
+       "<title>model</title>\n",
+       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-2366 1847.5,-2366 1847.5,4 -4,4\"/>\n",
+       "<!-- 0 -->\n",
+       "<g id=\"node1\" class=\"node\">\n",
+       "<title>0</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"878,-2362 754,-2362 754,-2330 878,-2330 878,-2362\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"754,-2330 754,-2362 824,-2362 824,-2330 754,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"759\" y=\"-2349\" font-family=\"Linux libertine\" font-size=\"10.00\">input&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"770.5\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"824,-2330 824,-2362 878,-2362 878,-2330 824,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"829\" y=\"-2343.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730)</text>\n",
+       "</g>\n",
+       "<!-- 1 -->\n",
+       "<g id=\"node2\" class=\"node\">\n",
+       "<title>1</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"342.5,-2294 167.5,-2294 167.5,-2252 342.5,-2252 342.5,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"168,-2252 168,-2294 238,-2294 238,-2252 168,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"173\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"184.5\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"238,-2273 238,-2294 286,-2294 286,-2273 238,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"248\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"286,-2273 286,-2294 343,-2294 343,-2273 286,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"291\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"238,-2252 238,-2273 286,-2273 286,-2252 238,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"243\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"286,-2252 286,-2273 343,-2273 343,-2252 286,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"297\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 0&#45;&gt;1 -->\n",
+       "<g id=\"edge1\" class=\"edge\">\n",
+       "<title>0&#45;&gt;1</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M753.82,-2337.13C656.76,-2324.85 468.43,-2301.01 352.98,-2286.4\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"353.1,-2282.89 342.74,-2285.1 352.22,-2289.83 353.1,-2282.89\"/>\n",
+       "</g>\n",
+       "<!-- 2 -->\n",
+       "<g id=\"node3\" class=\"node\">\n",
+       "<title>2</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"719.5,-2294 544.5,-2294 544.5,-2252 719.5,-2252 719.5,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"545,-2252 545,-2294 615,-2294 615,-2252 545,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"550\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"561.5\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2273 615,-2294 663,-2294 663,-2273 615,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"625\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2273 663,-2294 720,-2294 720,-2273 663,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2252 615,-2273 663,-2273 663,-2252 615,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"620\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2252 663,-2273 720,-2273 720,-2252 663,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "</g>\n",
+       "<!-- 0&#45;&gt;2 -->\n",
+       "<g id=\"edge2\" class=\"edge\">\n",
+       "<title>0&#45;&gt;2</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M776.98,-2329.94C752.73,-2320.58 721.02,-2308.35 693.45,-2297.71\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"694.48,-2294.36 683.89,-2294.02 691.96,-2300.89 694.48,-2294.36\"/>\n",
+       "</g>\n",
+       "<!-- 59 -->\n",
+       "<g id=\"node60\" class=\"node\">\n",
+       "<title>59</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1000.5,-2294 825.5,-2294 825.5,-2252 1000.5,-2252 1000.5,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"826,-2252 826,-2294 896,-2294 896,-2252 826,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"831\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"842.5\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"896,-2273 896,-2294 944,-2294 944,-2273 896,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"906\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-2273 944,-2294 1001,-2294 1001,-2273 944,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"949\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"896,-2252 896,-2273 944,-2273 944,-2252 896,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"901\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-2252 944,-2273 1001,-2273 1001,-2252 944,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"955\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 0&#45;&gt;59 -->\n",
+       "<g id=\"edge3\" class=\"edge\">\n",
+       "<title>0&#45;&gt;59</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M836.57,-2329.94C848.36,-2321.31 863.49,-2310.24 877.17,-2300.22\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"879.46,-2302.89 885.46,-2294.16 875.33,-2297.24 879.46,-2302.89\"/>\n",
+       "</g>\n",
+       "<!-- 60 -->\n",
+       "<g id=\"node61\" class=\"node\">\n",
+       "<title>60</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1457.5,-2294 1282.5,-2294 1282.5,-2252 1457.5,-2252 1457.5,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1283,-2252 1283,-2294 1353,-2294 1353,-2252 1283,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1288\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1299.5\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2273 1353,-2294 1401,-2294 1401,-2273 1353,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2273 1401,-2294 1458,-2294 1458,-2273 1401,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2252 1353,-2273 1401,-2273 1401,-2252 1353,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2252 1401,-2273 1458,-2273 1458,-2252 1401,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "</g>\n",
+       "<!-- 0&#45;&gt;60 -->\n",
+       "<g id=\"edge4\" class=\"edge\">\n",
+       "<title>0&#45;&gt;60</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M878.21,-2337.03C973.97,-2324.76 1158.31,-2301.13 1272.09,-2286.55\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1272.72,-2290 1282.2,-2285.25 1271.83,-2283.05 1272.72,-2290\"/>\n",
+       "</g>\n",
+       "<!-- 10 -->\n",
+       "<g id=\"node11\" class=\"node\">\n",
+       "<title>10</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"140,-2216 0,-2216 0,-2174 140,-2174 140,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"0,-2174 0,-2216 47,-2216 47,-2174 0,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"8\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"47,-2195 47,-2216 95,-2216 95,-2195 47,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"57\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"95,-2195 95,-2216 140,-2216 140,-2195 95,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"100\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"47,-2174 47,-2195 95,-2195 95,-2174 47,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"52\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"95,-2174 95,-2195 140,-2195 140,-2174 95,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"100\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 1&#45;&gt;10 -->\n",
+       "<g id=\"edge12\" class=\"edge\">\n",
+       "<title>1&#45;&gt;10</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M205.91,-2251.83C182.08,-2242.05 153.29,-2230.22 128.35,-2219.97\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"129.58,-2216.69 119,-2216.13 126.92,-2223.17 129.58,-2216.69\"/>\n",
+       "</g>\n",
+       "<!-- 3 -->\n",
+       "<g id=\"node4\" class=\"node\">\n",
+       "<title>3</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"526.5,-2216 351.5,-2216 351.5,-2174 526.5,-2174 526.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"352,-2174 352,-2216 422,-2216 422,-2174 352,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"357\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"368.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"422,-2195 422,-2216 470,-2216 470,-2195 422,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"432\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"470,-2195 470,-2216 527,-2216 527,-2195 470,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"475\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"422,-2174 422,-2195 470,-2195 470,-2174 422,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"427\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"470,-2174 470,-2195 527,-2195 527,-2174 470,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"478\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;3 -->\n",
+       "<g id=\"edge5\" class=\"edge\">\n",
+       "<title>2&#45;&gt;3</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M581.04,-2251.93C555.97,-2242.06 525.59,-2230.1 499.37,-2219.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"500.5,-2216.46 489.91,-2216.05 497.93,-2222.97 500.5,-2216.46\"/>\n",
+       "</g>\n",
+       "<!-- 5 -->\n",
+       "<g id=\"node6\" class=\"node\">\n",
+       "<title>5</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"912.5,-2216 737.5,-2216 737.5,-2174 912.5,-2174 912.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"738,-2174 738,-2216 808,-2216 808,-2174 738,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"743\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"754.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"808,-2195 808,-2216 856,-2216 856,-2195 808,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"818\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"856,-2195 856,-2216 913,-2216 913,-2195 856,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"861\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"808,-2174 808,-2195 856,-2195 856,-2174 808,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"813\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"856,-2174 856,-2195 913,-2195 913,-2174 856,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"861\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;5 -->\n",
+       "<g id=\"edge6\" class=\"edge\">\n",
+       "<title>2&#45;&gt;5</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M682.96,-2251.93C708.03,-2242.06 738.41,-2230.1 764.63,-2219.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"766.07,-2222.97 774.09,-2216.05 763.5,-2216.46 766.07,-2222.97\"/>\n",
+       "</g>\n",
+       "<!-- 7 -->\n",
+       "<g id=\"node8\" class=\"node\">\n",
+       "<title>7</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"333.5,-2216 158.5,-2216 158.5,-2174 333.5,-2174 333.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"159,-2174 159,-2216 229,-2216 229,-2174 159,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"164\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"175.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"229,-2195 229,-2216 277,-2216 277,-2195 229,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"239\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"277,-2195 277,-2216 334,-2216 334,-2195 277,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"282\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"229,-2174 229,-2195 277,-2195 277,-2174 229,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"234\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"277,-2174 277,-2195 334,-2195 334,-2174 277,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"291\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 0) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;7 -->\n",
+       "<g id=\"edge7\" class=\"edge\">\n",
+       "<title>2&#45;&gt;7</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M544.37,-2255.42C489.28,-2244.92 416.45,-2230.89 343.45,-2216.28\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"344.12,-2212.84 333.62,-2214.31 342.74,-2219.71 344.12,-2212.84\"/>\n",
+       "</g>\n",
+       "<!-- 8 -->\n",
+       "<g id=\"node9\" class=\"node\">\n",
+       "<title>8</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"719.5,-2216 544.5,-2216 544.5,-2174 719.5,-2174 719.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"545,-2174 545,-2216 615,-2216 615,-2174 545,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"550\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"561.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2195 615,-2216 663,-2216 663,-2195 615,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"625\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2195 663,-2216 720,-2216 720,-2195 663,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2174 615,-2195 663,-2195 663,-2174 615,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"620\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2174 663,-2195 720,-2195 720,-2174 663,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"674\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;8 -->\n",
+       "<g id=\"edge8\" class=\"edge\">\n",
+       "<title>2&#45;&gt;8</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2251.63C632,-2243.82 632,-2234.73 632,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2226.16 632,-2216.16 628.5,-2226.16 635.5,-2226.16\"/>\n",
+       "</g>\n",
+       "<!-- 4 -->\n",
+       "<g id=\"node5\" class=\"node\">\n",
+       "<title>4</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"519.5,-2138 358.5,-2138 358.5,-2096 519.5,-2096 519.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"359,-2096 359,-2138 409,-2138 409,-2096 359,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"364\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"365.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-2117 409,-2138 457,-2138 457,-2117 409,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"419\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"457,-2117 457,-2138 520,-2138 520,-2117 457,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"468\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-2096 409,-2117 457,-2117 457,-2096 409,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"414\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"457,-2096 457,-2117 520,-2117 520,-2096 457,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"462\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "</g>\n",
+       "<!-- 3&#45;&gt;4 -->\n",
+       "<g id=\"edge9\" class=\"edge\">\n",
+       "<title>3&#45;&gt;4</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M439,-2173.63C439,-2165.82 439,-2156.73 439,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"442.5,-2148.16 439,-2138.16 435.5,-2148.16 442.5,-2148.16\"/>\n",
+       "</g>\n",
+       "<!-- 14 -->\n",
+       "<g id=\"node15\" class=\"node\">\n",
+       "<title>14</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"518,-2060 360,-2060 360,-2018 518,-2018 518,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"360,-2018 360,-2060 407,-2060 407,-2018 360,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"368\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"365\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"407,-2039 407,-2060 455,-2060 455,-2039 407,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"417\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"455,-2039 455,-2060 518,-2060 518,-2039 455,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"460\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"407,-2018 407,-2039 455,-2039 455,-2018 407,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"412\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"455,-2018 455,-2039 518,-2039 518,-2018 455,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"460\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 4&#45;&gt;14 -->\n",
+       "<g id=\"edge17\" class=\"edge\">\n",
+       "<title>4&#45;&gt;14</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M439,-2095.63C439,-2087.82 439,-2078.73 439,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"442.5,-2070.16 439,-2060.16 435.5,-2070.16 442.5,-2070.16\"/>\n",
+       "</g>\n",
+       "<!-- 6 -->\n",
+       "<g id=\"node7\" class=\"node\">\n",
+       "<title>6</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"908.5,-2138 741.5,-2138 741.5,-2096 908.5,-2096 908.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"742,-2096 742,-2138 792,-2138 792,-2096 742,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"747\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"748.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"792,-2117 792,-2138 840,-2138 840,-2117 792,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"802\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"840,-2117 840,-2138 909,-2138 909,-2117 840,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"851\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"792,-2096 792,-2117 840,-2117 840,-2096 792,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"797\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"840,-2096 840,-2117 909,-2117 909,-2096 840,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"845\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "</g>\n",
+       "<!-- 5&#45;&gt;6 -->\n",
+       "<g id=\"edge10\" class=\"edge\">\n",
+       "<title>5&#45;&gt;6</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M825,-2173.63C825,-2165.82 825,-2156.73 825,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"828.5,-2148.16 825,-2138.16 821.5,-2148.16 828.5,-2148.16\"/>\n",
+       "</g>\n",
+       "<!-- 26 -->\n",
+       "<g id=\"node27\" class=\"node\">\n",
+       "<title>26</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"907,-2060 743,-2060 743,-2018 907,-2018 907,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"743,-2018 743,-2060 790,-2060 790,-2018 743,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"751\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"748\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"790,-2039 790,-2060 838,-2060 838,-2039 790,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"800\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"838,-2039 838,-2060 907,-2060 907,-2039 838,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"843\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"790,-2018 790,-2039 838,-2039 838,-2018 790,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"795\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"838,-2018 838,-2039 907,-2039 907,-2018 838,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"843\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 6&#45;&gt;26 -->\n",
+       "<g id=\"edge31\" class=\"edge\">\n",
+       "<title>6&#45;&gt;26</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M825,-2095.63C825,-2087.82 825,-2078.73 825,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"828.5,-2070.16 825,-2060.16 821.5,-2070.16 828.5,-2070.16\"/>\n",
+       "</g>\n",
+       "<!-- 9 -->\n",
+       "<g id=\"node10\" class=\"node\">\n",
+       "<title>9</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"704.5,-2138 555.5,-2138 555.5,-2096 704.5,-2096 704.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"556,-2096 556,-2138 606,-2138 606,-2096 556,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"562.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-2117 606,-2138 654,-2138 654,-2117 606,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"616\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-2117 654,-2138 705,-2138 705,-2117 654,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"662\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-2096 606,-2117 654,-2117 654,-2096 606,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"611\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-2096 654,-2117 705,-2117 705,-2096 654,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 8&#45;&gt;9 -->\n",
+       "<g id=\"edge11\" class=\"edge\">\n",
+       "<title>8&#45;&gt;9</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M631.46,-2173.63C631.26,-2165.82 631.02,-2156.73 630.79,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"634.29,-2148.06 630.53,-2138.16 627.29,-2148.25 634.29,-2148.06\"/>\n",
+       "</g>\n",
+       "<!-- 22 -->\n",
+       "<g id=\"node23\" class=\"node\">\n",
+       "<title>22</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"702,-2060 556,-2060 556,-2018 702,-2018 702,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"556,-2018 556,-2060 603,-2060 603,-2018 556,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"564\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"603,-2039 603,-2060 651,-2060 651,-2039 603,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"613\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"651,-2039 651,-2060 702,-2060 702,-2039 651,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"656\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"603,-2018 603,-2039 651,-2039 651,-2018 603,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"608\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"651,-2018 651,-2039 702,-2039 702,-2018 651,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"656\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 9&#45;&gt;22 -->\n",
+       "<g id=\"edge26\" class=\"edge\">\n",
+       "<title>9&#45;&gt;22</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629.73,-2095.63C629.63,-2087.82 629.51,-2078.73 629.4,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.9,-2070.11 629.27,-2060.16 625.9,-2070.2 632.9,-2070.11\"/>\n",
+       "</g>\n",
+       "<!-- 11 -->\n",
+       "<g id=\"node12\" class=\"node\">\n",
+       "<title>11</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"190.5,-2138 47.5,-2138 47.5,-2096 190.5,-2096 190.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"48,-2096 48,-2138 98,-2138 98,-2096 48,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"53\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"54.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"98,-2117 98,-2138 146,-2138 146,-2117 98,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"108\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"146,-2117 146,-2138 191,-2138 191,-2117 146,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"151\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"98,-2096 98,-2117 146,-2117 146,-2096 98,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"103\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"146,-2096 146,-2117 191,-2117 191,-2096 146,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"151\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 10&#45;&gt;11 -->\n",
+       "<g id=\"edge13\" class=\"edge\">\n",
+       "<title>10&#45;&gt;11</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M83.13,-2173.63C88.45,-2165.39 94.68,-2155.72 100.44,-2146.78\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"103.53,-2148.46 106,-2138.16 97.64,-2144.67 103.53,-2148.46\"/>\n",
+       "</g>\n",
+       "<!-- 12 -->\n",
+       "<g id=\"node13\" class=\"node\">\n",
+       "<title>12</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"294,-2060 136,-2060 136,-2018 294,-2018 294,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"136,-2018 136,-2060 201,-2060 201,-2018 136,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"141\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"150\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"201,-2039 201,-2060 249,-2060 249,-2039 201,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"211\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"249,-2039 249,-2060 294,-2060 294,-2039 249,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"254\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"201,-2018 201,-2039 249,-2039 249,-2018 201,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"206\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"249,-2018 249,-2039 294,-2039 294,-2018 249,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"254\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 11&#45;&gt;12 -->\n",
+       "<g id=\"edge14\" class=\"edge\">\n",
+       "<title>11&#45;&gt;12</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M144.47,-2095.83C155.76,-2086.9 169.19,-2076.27 181.29,-2066.69\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"183.74,-2069.21 189.41,-2060.26 179.39,-2063.72 183.74,-2069.21\"/>\n",
+       "</g>\n",
+       "<!-- 13 -->\n",
+       "<g id=\"node14\" class=\"node\">\n",
+       "<title>13</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"321,-1514 181,-1514 181,-1472 321,-1472 321,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"181,-1472 181,-1514 228,-1514 228,-1472 181,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"193\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"186\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"228,-1493 228,-1514 276,-1514 276,-1493 228,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"238\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"276,-1493 276,-1514 321,-1514 321,-1493 276,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"281\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"228,-1472 228,-1493 276,-1493 276,-1472 228,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"233\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"276,-1472 276,-1493 321,-1493 321,-1472 276,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"281\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 12&#45;&gt;13 -->\n",
+       "<g id=\"edge15\" class=\"edge\">\n",
+       "<title>12&#45;&gt;13</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M220.86,-2017.77C228.8,-1988.52 242,-1932.59 242,-1884 242,-1884 242,-1884 242,-1648 242,-1605.12 245.52,-1555.63 248.16,-1524.58\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"251.66,-1524.67 249.05,-1514.4 244.69,-1524.06 251.66,-1524.67\"/>\n",
+       "</g>\n",
+       "<!-- 43 -->\n",
+       "<g id=\"node44\" class=\"node\">\n",
+       "<title>43</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"651.5,-1202 424.5,-1202 424.5,-1160 651.5,-1160 651.5,-1202\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"425,-1160 425,-1202 472,-1202 472,-1160 425,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"440.5\" y=\"-1184\" font-family=\"Linux libertine\" font-size=\"10.00\">cat</text>\n",
+       "<text text-anchor=\"start\" x=\"430\" y=\"-1173\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"472,-1181 472,-1202 520,-1202 520,-1181 472,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"482\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"520,-1181 520,-1202 652,-1202 652,-1181 520,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"525\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10), 2 x (1, 32), (1, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"472,-1160 472,-1181 520,-1181 520,-1160 472,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"477\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"520,-1160 520,-1181 652,-1181 652,-1160 520,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"568.5\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "</g>\n",
+       "<!-- 13&#45;&gt;43 -->\n",
+       "<g id=\"edge16\" class=\"edge\">\n",
+       "<title>13&#45;&gt;43</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M255.59,-1471.65C267.39,-1423.91 303.25,-1303.81 378,-1238 393.83,-1224.07 413.33,-1213.54 433.21,-1205.62\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"434.55,-1208.85 442.67,-1202.04 432.08,-1202.3 434.55,-1208.85\"/>\n",
+       "</g>\n",
+       "<!-- 15 -->\n",
+       "<g id=\"node16\" class=\"node\">\n",
+       "<title>15</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"520.5,-1982 359.5,-1982 359.5,-1940 520.5,-1940 520.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"360,-1940 360,-1982 410,-1982 410,-1940 360,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"365\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"366.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"410,-1961 410,-1982 458,-1982 458,-1961 410,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"420\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"458,-1961 458,-1982 521,-1982 521,-1961 458,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"463\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"410,-1940 410,-1961 458,-1961 458,-1940 410,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"415\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"458,-1940 458,-1961 521,-1961 521,-1940 458,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"463\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 14&#45;&gt;15 -->\n",
+       "<g id=\"edge18\" class=\"edge\">\n",
+       "<title>14&#45;&gt;15</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M439.27,-2017.63C439.37,-2009.82 439.49,-2000.73 439.6,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"443.1,-1992.2 439.73,-1982.16 436.1,-1992.11 443.1,-1992.2\"/>\n",
+       "</g>\n",
+       "<!-- 16 -->\n",
+       "<g id=\"node17\" class=\"node\">\n",
+       "<title>16</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"529,-1904 353,-1904 353,-1862 529,-1862 529,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"353,-1862 353,-1904 418,-1904 418,-1862 353,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"358\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"367\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"418,-1883 418,-1904 466,-1904 466,-1883 418,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"428\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"466,-1883 466,-1904 529,-1904 529,-1883 466,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"471\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"418,-1862 418,-1883 466,-1883 466,-1862 418,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"423\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"466,-1862 466,-1883 529,-1883 529,-1862 466,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"471\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 15&#45;&gt;16 -->\n",
+       "<g id=\"edge19\" class=\"edge\">\n",
+       "<title>15&#45;&gt;16</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M440.27,-1939.63C440.37,-1931.82 440.49,-1922.73 440.6,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"444.1,-1914.2 440.73,-1904.16 437.1,-1914.11 444.1,-1914.2\"/>\n",
+       "</g>\n",
+       "<!-- 17 -->\n",
+       "<g id=\"node18\" class=\"node\">\n",
+       "<title>17</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"521,-1826 363,-1826 363,-1784 521,-1784 521,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"363,-1784 363,-1826 410,-1826 410,-1784 363,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"375\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"368\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"410,-1805 410,-1826 458,-1826 458,-1805 410,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"420\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"458,-1805 458,-1826 521,-1826 521,-1805 458,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"463\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"410,-1784 410,-1805 458,-1805 458,-1784 410,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"415\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"458,-1784 458,-1805 521,-1805 521,-1784 458,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"463\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 16&#45;&gt;17 -->\n",
+       "<g id=\"edge20\" class=\"edge\">\n",
+       "<title>16&#45;&gt;17</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M441.27,-1861.63C441.37,-1853.82 441.49,-1844.73 441.6,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"445.1,-1836.2 441.73,-1826.16 438.1,-1836.11 445.1,-1836.2\"/>\n",
+       "</g>\n",
+       "<!-- 18 -->\n",
+       "<g id=\"node19\" class=\"node\">\n",
+       "<title>18</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"523,-1748 365,-1748 365,-1706 523,-1706 523,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"365,-1706 365,-1748 412,-1748 412,-1706 365,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"373\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"370\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"412,-1727 412,-1748 460,-1748 460,-1727 412,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"422\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-1727 460,-1748 523,-1748 523,-1727 460,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"412,-1706 412,-1727 460,-1727 460,-1706 412,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"417\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-1706 460,-1727 523,-1727 523,-1706 460,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 17&#45;&gt;18 -->\n",
+       "<g id=\"edge21\" class=\"edge\">\n",
+       "<title>17&#45;&gt;18</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M442.54,-1783.63C442.74,-1775.82 442.98,-1766.73 443.21,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"446.71,-1758.25 443.47,-1748.16 439.71,-1758.06 446.71,-1758.25\"/>\n",
+       "</g>\n",
+       "<!-- 19 -->\n",
+       "<g id=\"node20\" class=\"node\">\n",
+       "<title>19</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"525.5,-1670 364.5,-1670 364.5,-1628 525.5,-1628 525.5,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"365,-1628 365,-1670 415,-1670 415,-1628 365,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"370\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"371.5\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"415,-1649 415,-1670 463,-1670 463,-1649 415,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"425\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"463,-1649 463,-1670 526,-1670 526,-1649 463,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"468\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"415,-1628 415,-1649 463,-1649 463,-1628 415,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"420\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"463,-1628 463,-1649 526,-1649 526,-1628 463,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"468\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 18&#45;&gt;19 -->\n",
+       "<g id=\"edge22\" class=\"edge\">\n",
+       "<title>18&#45;&gt;19</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M444.27,-1705.63C444.37,-1697.82 444.49,-1688.73 444.6,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"448.1,-1680.2 444.73,-1670.16 441.1,-1680.11 448.1,-1680.2\"/>\n",
+       "</g>\n",
+       "<!-- 20 -->\n",
+       "<g id=\"node21\" class=\"node\">\n",
+       "<title>20</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"533,-1592 357,-1592 357,-1550 533,-1550 533,-1592\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"357,-1550 357,-1592 422,-1592 422,-1550 357,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"362\" y=\"-1574\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"371\" y=\"-1563\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"422,-1571 422,-1592 470,-1592 470,-1571 422,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"432\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"470,-1571 470,-1592 533,-1592 533,-1571 470,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"475\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"422,-1550 422,-1571 470,-1571 470,-1550 422,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"427\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"470,-1550 470,-1571 533,-1571 533,-1550 470,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"475\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 19&#45;&gt;20 -->\n",
+       "<g id=\"edge23\" class=\"edge\">\n",
+       "<title>19&#45;&gt;20</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M445,-1627.63C445,-1619.82 445,-1610.73 445,-1602.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"448.5,-1602.16 445,-1592.16 441.5,-1602.16 448.5,-1602.16\"/>\n",
+       "</g>\n",
+       "<!-- 21 -->\n",
+       "<g id=\"node22\" class=\"node\">\n",
+       "<title>21</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"524,-1514 366,-1514 366,-1472 524,-1472 524,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"366,-1472 366,-1514 413,-1514 413,-1472 366,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"378\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"371\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"413,-1493 413,-1514 461,-1514 461,-1493 413,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"423\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"461,-1493 461,-1514 524,-1514 524,-1493 461,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"466\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"413,-1472 413,-1493 461,-1493 461,-1472 413,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"418\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"461,-1472 461,-1493 524,-1493 524,-1472 461,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"466\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 20&#45;&gt;21 -->\n",
+       "<g id=\"edge24\" class=\"edge\">\n",
+       "<title>20&#45;&gt;21</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M445,-1549.63C445,-1541.82 445,-1532.73 445,-1524.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"448.5,-1524.16 445,-1514.16 441.5,-1524.16 448.5,-1524.16\"/>\n",
+       "</g>\n",
+       "<!-- 34 -->\n",
+       "<g id=\"node35\" class=\"node\">\n",
+       "<title>34</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"526.5,-1436 363.5,-1436 363.5,-1394 526.5,-1394 526.5,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"364,-1394 364,-1436 416,-1436 416,-1394 364,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"369\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"371.5\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"416,-1415 416,-1436 464,-1436 464,-1415 416,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"426\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"464,-1415 464,-1436 527,-1436 527,-1415 464,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"469\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"416,-1394 416,-1415 464,-1415 464,-1394 416,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"421\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"464,-1394 464,-1415 527,-1415 527,-1394 464,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"469\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "</g>\n",
+       "<!-- 21&#45;&gt;34 -->\n",
+       "<g id=\"edge25\" class=\"edge\">\n",
+       "<title>21&#45;&gt;34</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M445,-1471.63C445,-1463.82 445,-1454.73 445,-1446.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"448.5,-1446.16 445,-1436.16 441.5,-1446.16 448.5,-1446.16\"/>\n",
+       "</g>\n",
+       "<!-- 23 -->\n",
+       "<g id=\"node24\" class=\"node\">\n",
+       "<title>23</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"703.5,-1982 554.5,-1982 554.5,-1940 703.5,-1940 703.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"555,-1940 555,-1982 605,-1982 605,-1940 555,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"560\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"561.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"605,-1961 605,-1982 653,-1982 653,-1961 605,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"615\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"653,-1961 653,-1982 704,-1982 704,-1961 653,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"658\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"605,-1940 605,-1961 653,-1961 653,-1940 605,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"610\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"653,-1940 653,-1961 704,-1961 704,-1940 653,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"658\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 22&#45;&gt;23 -->\n",
+       "<g id=\"edge27\" class=\"edge\">\n",
+       "<title>22&#45;&gt;23</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629,-2017.63C629,-2009.82 629,-2000.73 629,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.5,-1992.16 629,-1982.16 625.5,-1992.16 632.5,-1992.16\"/>\n",
+       "</g>\n",
+       "<!-- 24 -->\n",
+       "<g id=\"node25\" class=\"node\">\n",
+       "<title>24</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"711,-1904 547,-1904 547,-1862 711,-1862 711,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"547,-1862 547,-1904 612,-1904 612,-1862 547,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"552\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"612,-1883 612,-1904 660,-1904 660,-1883 612,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"622\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"660,-1883 660,-1904 711,-1904 711,-1883 660,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"665\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"612,-1862 612,-1883 660,-1883 660,-1862 612,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"617\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"660,-1862 660,-1883 711,-1883 711,-1862 660,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"665\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 23&#45;&gt;24 -->\n",
+       "<g id=\"edge28\" class=\"edge\">\n",
+       "<title>23&#45;&gt;24</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629,-1939.63C629,-1931.82 629,-1922.73 629,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.5,-1914.16 629,-1904.16 625.5,-1914.16 632.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 25 -->\n",
+       "<g id=\"node26\" class=\"node\">\n",
+       "<title>25</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"702,-1826 556,-1826 556,-1784 702,-1784 702,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"556,-1784 556,-1826 603,-1826 603,-1784 556,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"603,-1805 603,-1826 651,-1826 651,-1805 603,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"613\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"651,-1805 651,-1826 702,-1826 702,-1805 651,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"656\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"603,-1784 603,-1805 651,-1805 651,-1784 603,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"608\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"651,-1784 651,-1805 702,-1805 702,-1784 651,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"656\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 24&#45;&gt;25 -->\n",
+       "<g id=\"edge29\" class=\"edge\">\n",
+       "<title>24&#45;&gt;25</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629,-1861.63C629,-1853.82 629,-1844.73 629,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.5,-1836.16 629,-1826.16 625.5,-1836.16 632.5,-1836.16\"/>\n",
+       "</g>\n",
+       "<!-- 37 -->\n",
+       "<g id=\"node38\" class=\"node\">\n",
+       "<title>37</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"704.5,-1748 553.5,-1748 553.5,-1706 704.5,-1706 704.5,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"554,-1706 554,-1748 606,-1748 606,-1706 554,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"559\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"561.5\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-1727 606,-1748 654,-1748 654,-1727 606,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"616\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-1727 654,-1748 705,-1748 705,-1727 654,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-1706 606,-1727 654,-1727 654,-1706 606,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"611\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-1706 654,-1727 705,-1727 705,-1706 654,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "</g>\n",
+       "<!-- 25&#45;&gt;37 -->\n",
+       "<g id=\"edge30\" class=\"edge\">\n",
+       "<title>25&#45;&gt;37</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629,-1783.63C629,-1775.82 629,-1766.73 629,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.5,-1758.16 629,-1748.16 625.5,-1758.16 632.5,-1758.16\"/>\n",
+       "</g>\n",
+       "<!-- 27 -->\n",
+       "<g id=\"node28\" class=\"node\">\n",
+       "<title>27</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"907.5,-1982 740.5,-1982 740.5,-1940 907.5,-1940 907.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"741,-1940 741,-1982 791,-1982 791,-1940 741,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"746\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"747.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-1961 791,-1982 839,-1982 839,-1961 791,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"801\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-1961 839,-1982 908,-1982 908,-1961 839,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-1940 791,-1961 839,-1961 839,-1940 791,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"796\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-1940 839,-1961 908,-1961 908,-1940 839,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 26&#45;&gt;27 -->\n",
+       "<g id=\"edge32\" class=\"edge\">\n",
+       "<title>26&#45;&gt;27</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M824.73,-2017.63C824.63,-2009.82 824.51,-2000.73 824.4,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"827.9,-1992.11 824.27,-1982.16 820.9,-1992.2 827.9,-1992.11\"/>\n",
+       "</g>\n",
+       "<!-- 28 -->\n",
+       "<g id=\"node29\" class=\"node\">\n",
+       "<title>28</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"913,-1904 731,-1904 731,-1862 913,-1862 913,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"731,-1862 731,-1904 796,-1904 796,-1862 731,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"736\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"745\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"796,-1883 796,-1904 844,-1904 844,-1883 796,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"806\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"844,-1883 844,-1904 913,-1904 913,-1883 844,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"849\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"796,-1862 796,-1883 844,-1883 844,-1862 796,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"801\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"844,-1862 844,-1883 913,-1883 913,-1862 844,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"849\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 27&#45;&gt;28 -->\n",
+       "<g id=\"edge33\" class=\"edge\">\n",
+       "<title>27&#45;&gt;28</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M823.46,-1939.63C823.26,-1931.82 823.02,-1922.73 822.79,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"826.29,-1914.06 822.53,-1904.16 819.29,-1914.25 826.29,-1914.06\"/>\n",
+       "</g>\n",
+       "<!-- 29 -->\n",
+       "<g id=\"node30\" class=\"node\">\n",
+       "<title>29</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"903,-1826 739,-1826 739,-1784 903,-1784 903,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"739,-1784 739,-1826 786,-1826 786,-1784 739,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"751\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"744\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"786,-1805 786,-1826 834,-1826 834,-1805 786,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"796\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"834,-1805 834,-1826 903,-1826 903,-1805 834,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"839\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"786,-1784 786,-1805 834,-1805 834,-1784 786,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"791\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"834,-1784 834,-1805 903,-1805 903,-1784 834,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"839\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 28&#45;&gt;29 -->\n",
+       "<g id=\"edge34\" class=\"edge\">\n",
+       "<title>28&#45;&gt;29</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M821.73,-1861.63C821.63,-1853.82 821.51,-1844.73 821.4,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"824.9,-1836.11 821.27,-1826.16 817.9,-1836.2 824.9,-1836.11\"/>\n",
+       "</g>\n",
+       "<!-- 30 -->\n",
+       "<g id=\"node31\" class=\"node\">\n",
+       "<title>30</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"900,-1748 736,-1748 736,-1706 900,-1706 900,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"736,-1706 736,-1748 783,-1748 783,-1706 736,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"744\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"741\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"783,-1727 783,-1748 831,-1748 831,-1727 783,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"793\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"831,-1727 831,-1748 900,-1748 900,-1727 831,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"836\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"783,-1706 783,-1727 831,-1727 831,-1706 783,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"788\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"831,-1706 831,-1727 900,-1727 900,-1706 831,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"836\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 29&#45;&gt;30 -->\n",
+       "<g id=\"edge35\" class=\"edge\">\n",
+       "<title>29&#45;&gt;30</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M820.2,-1783.63C819.89,-1775.82 819.53,-1766.73 819.19,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"822.69,-1758.01 818.8,-1748.16 815.69,-1758.29 822.69,-1758.01\"/>\n",
+       "</g>\n",
+       "<!-- 31 -->\n",
+       "<g id=\"node32\" class=\"node\">\n",
+       "<title>31</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"901.5,-1670 734.5,-1670 734.5,-1628 901.5,-1628 901.5,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"735,-1628 735,-1670 785,-1670 785,-1628 735,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"740\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"741.5\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"785,-1649 785,-1670 833,-1670 833,-1649 785,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"795\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"833,-1649 833,-1670 902,-1670 902,-1649 833,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"838\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"785,-1628 785,-1649 833,-1649 833,-1628 785,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"790\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"833,-1628 833,-1649 902,-1649 902,-1628 833,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"838\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 30&#45;&gt;31 -->\n",
+       "<g id=\"edge36\" class=\"edge\">\n",
+       "<title>30&#45;&gt;31</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M818,-1705.63C818,-1697.82 818,-1688.73 818,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"821.5,-1680.16 818,-1670.16 814.5,-1680.16 821.5,-1680.16\"/>\n",
+       "</g>\n",
+       "<!-- 32 -->\n",
+       "<g id=\"node33\" class=\"node\">\n",
+       "<title>32</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"908,-1592 726,-1592 726,-1550 908,-1550 908,-1592\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"726,-1550 726,-1592 791,-1592 791,-1550 726,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"731\" y=\"-1574\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"740\" y=\"-1563\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-1571 791,-1592 839,-1592 839,-1571 791,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"801\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-1571 839,-1592 908,-1592 908,-1571 839,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-1550 791,-1571 839,-1571 839,-1550 791,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"796\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-1550 839,-1571 908,-1571 908,-1550 839,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 31&#45;&gt;32 -->\n",
+       "<g id=\"edge37\" class=\"edge\">\n",
+       "<title>31&#45;&gt;32</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M817.73,-1627.63C817.63,-1619.82 817.51,-1610.73 817.4,-1602.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"820.9,-1602.11 817.27,-1592.16 813.9,-1602.2 820.9,-1602.11\"/>\n",
+       "</g>\n",
+       "<!-- 33 -->\n",
+       "<g id=\"node34\" class=\"node\">\n",
+       "<title>33</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"899,-1514 735,-1514 735,-1472 899,-1472 899,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"735,-1472 735,-1514 782,-1514 782,-1472 735,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"747\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"740\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"782,-1493 782,-1514 830,-1514 830,-1493 782,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"792\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"830,-1493 830,-1514 899,-1514 899,-1493 830,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"835\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"782,-1472 782,-1493 830,-1493 830,-1472 782,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"787\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"830,-1472 830,-1493 899,-1493 899,-1472 830,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"835\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 32&#45;&gt;33 -->\n",
+       "<g id=\"edge38\" class=\"edge\">\n",
+       "<title>32&#45;&gt;33</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M817,-1549.63C817,-1541.82 817,-1532.73 817,-1524.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"820.5,-1524.16 817,-1514.16 813.5,-1524.16 820.5,-1524.16\"/>\n",
+       "</g>\n",
+       "<!-- 40 -->\n",
+       "<g id=\"node41\" class=\"node\">\n",
+       "<title>40</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"901.5,-1436 732.5,-1436 732.5,-1394 901.5,-1394 901.5,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"733,-1394 733,-1436 785,-1436 785,-1394 733,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"738\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"740.5\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"785,-1415 785,-1436 833,-1436 833,-1415 785,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"795\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"833,-1415 833,-1436 902,-1436 902,-1415 833,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"838\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"785,-1394 785,-1415 833,-1415 833,-1394 785,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"790\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"833,-1394 833,-1415 902,-1415 902,-1394 833,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"838\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "</g>\n",
+       "<!-- 33&#45;&gt;40 -->\n",
+       "<g id=\"edge39\" class=\"edge\">\n",
+       "<title>33&#45;&gt;40</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M817,-1471.63C817,-1463.82 817,-1454.73 817,-1446.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"820.5,-1446.16 817,-1436.16 813.5,-1446.16 820.5,-1446.16\"/>\n",
+       "</g>\n",
+       "<!-- 35 -->\n",
+       "<g id=\"node36\" class=\"node\">\n",
+       "<title>35</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"536,-1358 354,-1358 354,-1316 536,-1316 536,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"354,-1316 354,-1358 425,-1358 425,-1316 354,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"359\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"371\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"425,-1337 425,-1358 473,-1358 473,-1337 425,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"435\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"473,-1337 473,-1358 536,-1358 536,-1337 473,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"478\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"425,-1316 425,-1337 473,-1337 473,-1316 425,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"430\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"473,-1316 473,-1337 536,-1337 536,-1316 473,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"481\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 34&#45;&gt;35 -->\n",
+       "<g id=\"edge40\" class=\"edge\">\n",
+       "<title>34&#45;&gt;35</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M445,-1393.63C445,-1385.82 445,-1376.73 445,-1368.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"448.5,-1368.16 445,-1358.16 441.5,-1368.16 448.5,-1368.16\"/>\n",
+       "</g>\n",
+       "<!-- 36 -->\n",
+       "<g id=\"node37\" class=\"node\">\n",
+       "<title>36</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"543,-1280 387,-1280 387,-1238 543,-1238 543,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"387,-1238 387,-1280 438,-1280 438,-1238 387,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"392\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"394\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"438,-1259 438,-1280 486,-1280 486,-1259 438,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"448\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"486,-1259 486,-1280 543,-1280 543,-1259 486,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"491\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"438,-1238 438,-1259 486,-1259 486,-1238 438,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"443\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"486,-1238 486,-1259 543,-1259 543,-1238 486,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"497\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 35&#45;&gt;36 -->\n",
+       "<g id=\"edge41\" class=\"edge\">\n",
+       "<title>35&#45;&gt;36</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M450.36,-1315.63C452.44,-1307.73 454.86,-1298.53 457.13,-1289.9\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"460.53,-1290.72 459.69,-1280.16 453.77,-1288.94 460.53,-1290.72\"/>\n",
+       "</g>\n",
+       "<!-- 36&#45;&gt;43 -->\n",
+       "<g id=\"edge42\" class=\"edge\">\n",
+       "<title>36&#45;&gt;43</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M484.56,-1237.63C492.89,-1228.96 502.74,-1218.71 511.69,-1209.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"514.23,-1211.8 518.64,-1202.16 509.18,-1206.95 514.23,-1211.8\"/>\n",
+       "</g>\n",
+       "<!-- 38 -->\n",
+       "<g id=\"node39\" class=\"node\">\n",
+       "<title>38</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"714,-1670 544,-1670 544,-1628 714,-1628 714,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"544,-1628 544,-1670 615,-1670 615,-1628 544,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"549\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-1649 615,-1670 663,-1670 663,-1649 615,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"625\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-1649 663,-1670 714,-1670 714,-1649 663,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-1628 615,-1649 663,-1649 663,-1628 615,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"620\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-1628 663,-1649 714,-1649 714,-1628 663,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "</g>\n",
+       "<!-- 37&#45;&gt;38 -->\n",
+       "<g id=\"edge43\" class=\"edge\">\n",
+       "<title>37&#45;&gt;38</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629,-1705.63C629,-1697.82 629,-1688.73 629,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.5,-1680.16 629,-1670.16 625.5,-1680.16 632.5,-1680.16\"/>\n",
+       "</g>\n",
+       "<!-- 39 -->\n",
+       "<g id=\"node40\" class=\"node\">\n",
+       "<title>39</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"704,-1358 554,-1358 554,-1316 704,-1316 704,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"554,-1316 554,-1358 605,-1358 605,-1316 554,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"559\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"605,-1337 605,-1358 653,-1358 653,-1337 605,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"615\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"653,-1337 653,-1358 704,-1358 704,-1337 653,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"658\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"605,-1316 605,-1337 653,-1337 653,-1316 605,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"610\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"653,-1316 653,-1337 704,-1337 704,-1316 653,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"664\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3) </text>\n",
+       "</g>\n",
+       "<!-- 38&#45;&gt;39 -->\n",
+       "<g id=\"edge44\" class=\"edge\">\n",
+       "<title>38&#45;&gt;39</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M629,-1627.66C629,-1574.43 629,-1430.84 629,-1368.06\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"632.5,-1368.01 629,-1358.01 625.5,-1368.01 632.5,-1368.01\"/>\n",
+       "</g>\n",
+       "<!-- 39&#45;&gt;43 -->\n",
+       "<g id=\"edge45\" class=\"edge\">\n",
+       "<title>39&#45;&gt;43</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M617.1,-1315.87C601.3,-1289.12 573.27,-1241.69 555.17,-1211.06\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"558,-1208.97 549.9,-1202.14 551.97,-1212.53 558,-1208.97\"/>\n",
+       "</g>\n",
+       "<!-- 41 -->\n",
+       "<g id=\"node42\" class=\"node\">\n",
+       "<title>41</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"910,-1358 722,-1358 722,-1316 910,-1316 910,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"722,-1316 722,-1358 793,-1358 793,-1316 722,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"727\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"739\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"793,-1337 793,-1358 841,-1358 841,-1337 793,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"803\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"841,-1337 841,-1358 910,-1358 910,-1337 841,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"846\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"793,-1316 793,-1337 841,-1337 841,-1316 793,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"798\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"841,-1316 841,-1337 910,-1337 910,-1316 841,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"852\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 40&#45;&gt;41 -->\n",
+       "<g id=\"edge46\" class=\"edge\">\n",
+       "<title>40&#45;&gt;41</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M816.73,-1393.63C816.63,-1385.82 816.51,-1376.73 816.4,-1368.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"819.9,-1368.11 816.27,-1358.16 812.9,-1368.2 819.9,-1368.11\"/>\n",
+       "</g>\n",
+       "<!-- 42 -->\n",
+       "<g id=\"node43\" class=\"node\">\n",
+       "<title>42</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"834,-1280 678,-1280 678,-1238 834,-1238 834,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"678,-1238 678,-1280 729,-1280 729,-1238 678,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"683\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"685\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"729,-1259 729,-1280 777,-1280 777,-1259 729,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"739\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"777,-1259 777,-1280 834,-1280 834,-1259 777,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"782\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"729,-1238 729,-1259 777,-1259 777,-1238 729,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"734\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"777,-1238 777,-1259 834,-1259 834,-1238 777,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"788\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 41&#45;&gt;42 -->\n",
+       "<g id=\"edge47\" class=\"edge\">\n",
+       "<title>41&#45;&gt;42</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M799.92,-1315.63C793.28,-1307.22 785.46,-1297.32 778.28,-1288.22\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"780.86,-1285.84 771.92,-1280.16 775.36,-1290.18 780.86,-1285.84\"/>\n",
+       "</g>\n",
+       "<!-- 42&#45;&gt;43 -->\n",
+       "<g id=\"edge48\" class=\"edge\">\n",
+       "<title>42&#45;&gt;43</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M698.44,-1237.93C669.75,-1227.93 634.9,-1215.78 605.02,-1205.36\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"606.1,-1202.04 595.51,-1202.05 603.8,-1208.65 606.1,-1202.04\"/>\n",
+       "</g>\n",
+       "<!-- 44 -->\n",
+       "<g id=\"node45\" class=\"node\">\n",
+       "<title>44</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"611,-1124 465,-1124 465,-1082 611,-1082 611,-1124\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-1082 465,-1124 512,-1124 512,-1082 465,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-1106\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-1095\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-1103 512,-1124 560,-1124 560,-1103 512,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"522\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-1103 560,-1124 611,-1124 611,-1103 560,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-1082 512,-1103 560,-1103 560,-1082 512,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"517\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-1082 560,-1103 611,-1103 611,-1082 560,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 43&#45;&gt;44 -->\n",
+       "<g id=\"edge49\" class=\"edge\">\n",
+       "<title>43&#45;&gt;44</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-1159.63C538,-1151.82 538,-1142.73 538,-1134.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-1134.16 538,-1124.16 534.5,-1134.16 541.5,-1134.16\"/>\n",
+       "</g>\n",
+       "<!-- 45 -->\n",
+       "<g id=\"node46\" class=\"node\">\n",
+       "<title>45</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"620,-1046 456,-1046 456,-1004 620,-1004 620,-1046\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"456,-1004 456,-1046 521,-1046 521,-1004 456,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"461\" y=\"-1028\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-1017\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"521,-1025 521,-1046 569,-1046 569,-1025 521,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"531\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"569,-1025 569,-1046 620,-1046 620,-1025 569,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"574\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"521,-1004 521,-1025 569,-1025 569,-1004 521,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"526\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"569,-1004 569,-1025 620,-1025 620,-1004 569,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"574\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 44&#45;&gt;45 -->\n",
+       "<g id=\"edge50\" class=\"edge\">\n",
+       "<title>44&#45;&gt;45</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-1081.63C538,-1073.82 538,-1064.73 538,-1056.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-1056.16 538,-1046.16 534.5,-1056.16 541.5,-1056.16\"/>\n",
+       "</g>\n",
+       "<!-- 46 -->\n",
+       "<g id=\"node47\" class=\"node\">\n",
+       "<title>46</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"611,-968 465,-968 465,-926 611,-926 611,-968\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-926 465,-968 512,-968 512,-926 465,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"477\" y=\"-950\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-939\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-947 512,-968 560,-968 560,-947 512,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"522\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-947 560,-968 611,-968 611,-947 560,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-926 512,-947 560,-947 560,-926 512,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"517\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-926 560,-947 611,-947 611,-926 560,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 45&#45;&gt;46 -->\n",
+       "<g id=\"edge51\" class=\"edge\">\n",
+       "<title>45&#45;&gt;46</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-1003.63C538,-995.82 538,-986.73 538,-978.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-978.16 538,-968.16 534.5,-978.16 541.5,-978.16\"/>\n",
+       "</g>\n",
+       "<!-- 47 -->\n",
+       "<g id=\"node48\" class=\"node\">\n",
+       "<title>47</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"612.5,-890 463.5,-890 463.5,-848 612.5,-848 612.5,-890\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"464,-848 464,-890 514,-890 514,-848 464,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"469\" y=\"-872\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"470.5\" y=\"-861\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"514,-869 514,-890 562,-890 562,-869 514,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"524\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"562,-869 562,-890 613,-890 613,-869 562,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"567\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"514,-848 514,-869 562,-869 562,-848 514,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"519\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"562,-848 562,-869 613,-869 613,-848 562,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"567\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 46&#45;&gt;47 -->\n",
+       "<g id=\"edge52\" class=\"edge\">\n",
+       "<title>46&#45;&gt;47</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-925.63C538,-917.82 538,-908.73 538,-900.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-900.16 538,-890.16 534.5,-900.16 541.5,-900.16\"/>\n",
+       "</g>\n",
+       "<!-- 48 -->\n",
+       "<g id=\"node49\" class=\"node\">\n",
+       "<title>48</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"611,-812 465,-812 465,-770 611,-770 611,-812\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-770 465,-812 512,-812 512,-770 465,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-794\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-783\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-791 512,-812 560,-812 560,-791 512,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"522\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-791 560,-812 611,-812 611,-791 560,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-770 512,-791 560,-791 560,-770 512,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"517\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-770 560,-791 611,-791 611,-770 560,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 47&#45;&gt;48 -->\n",
+       "<g id=\"edge53\" class=\"edge\">\n",
+       "<title>47&#45;&gt;48</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-847.63C538,-839.82 538,-830.73 538,-822.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-822.16 538,-812.16 534.5,-822.16 541.5,-822.16\"/>\n",
+       "</g>\n",
+       "<!-- 49 -->\n",
+       "<g id=\"node50\" class=\"node\">\n",
+       "<title>49</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"620,-734 456,-734 456,-692 620,-692 620,-734\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"456,-692 456,-734 521,-734 521,-692 456,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"461\" y=\"-716\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-705\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"521,-713 521,-734 569,-734 569,-713 521,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"531\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"569,-713 569,-734 620,-734 620,-713 569,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"574\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"521,-692 521,-713 569,-713 569,-692 521,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"526\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"569,-692 569,-713 620,-713 620,-692 569,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"574\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 48&#45;&gt;49 -->\n",
+       "<g id=\"edge54\" class=\"edge\">\n",
+       "<title>48&#45;&gt;49</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-769.63C538,-761.82 538,-752.73 538,-744.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-744.16 538,-734.16 534.5,-744.16 541.5,-744.16\"/>\n",
+       "</g>\n",
+       "<!-- 50 -->\n",
+       "<g id=\"node51\" class=\"node\">\n",
+       "<title>50</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"611,-656 465,-656 465,-614 611,-614 611,-656\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-614 465,-656 512,-656 512,-614 465,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"477\" y=\"-638\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-627\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-635 512,-656 560,-656 560,-635 512,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"522\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-635 560,-656 611,-656 611,-635 560,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-614 512,-635 560,-635 560,-614 512,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"517\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-614 560,-635 611,-635 611,-614 560,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 49&#45;&gt;50 -->\n",
+       "<g id=\"edge55\" class=\"edge\">\n",
+       "<title>49&#45;&gt;50</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-691.63C538,-683.82 538,-674.73 538,-666.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-666.16 538,-656.16 534.5,-666.16 541.5,-666.16\"/>\n",
+       "</g>\n",
+       "<!-- 51 -->\n",
+       "<g id=\"node52\" class=\"node\">\n",
+       "<title>51</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"612.5,-578 463.5,-578 463.5,-536 612.5,-536 612.5,-578\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"464,-536 464,-578 514,-578 514,-536 464,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"469\" y=\"-560\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"470.5\" y=\"-549\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"514,-557 514,-578 562,-578 562,-557 514,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"524\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"562,-557 562,-578 613,-578 613,-557 562,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"567\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"514,-536 514,-557 562,-557 562,-536 514,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"519\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"562,-536 562,-557 613,-557 613,-536 562,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"567\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 50&#45;&gt;51 -->\n",
+       "<g id=\"edge56\" class=\"edge\">\n",
+       "<title>50&#45;&gt;51</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-613.63C538,-605.82 538,-596.73 538,-588.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-588.16 538,-578.16 534.5,-588.16 541.5,-588.16\"/>\n",
+       "</g>\n",
+       "<!-- 52 -->\n",
+       "<g id=\"node53\" class=\"node\">\n",
+       "<title>52</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"611,-500 465,-500 465,-458 611,-458 611,-500\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-458 465,-500 512,-500 512,-458 465,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-479 512,-500 560,-500 560,-479 512,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"522\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-479 560,-500 611,-500 611,-479 560,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"512,-458 512,-479 560,-479 560,-458 512,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"517\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-458 560,-479 611,-479 611,-458 560,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 51&#45;&gt;52 -->\n",
+       "<g id=\"edge57\" class=\"edge\">\n",
+       "<title>51&#45;&gt;52</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-535.63C538,-527.82 538,-518.73 538,-510.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-510.16 538,-500.16 534.5,-510.16 541.5,-510.16\"/>\n",
+       "</g>\n",
+       "<!-- 53 -->\n",
+       "<g id=\"node54\" class=\"node\">\n",
+       "<title>53</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"617,-422 459,-422 459,-380 617,-380 617,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"459,-380 459,-422 524,-422 524,-380 459,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"464\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"524,-401 524,-422 572,-422 572,-401 524,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"534\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"572,-401 572,-422 617,-422 617,-401 572,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"577\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"524,-380 524,-401 572,-401 572,-380 524,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"529\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"572,-380 572,-401 617,-401 617,-380 572,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"577\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 52&#45;&gt;53 -->\n",
+       "<g id=\"edge58\" class=\"edge\">\n",
+       "<title>52&#45;&gt;53</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-457.63C538,-449.82 538,-440.73 538,-432.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-432.16 538,-422.16 534.5,-432.16 541.5,-432.16\"/>\n",
+       "</g>\n",
+       "<!-- 54 -->\n",
+       "<g id=\"node55\" class=\"node\">\n",
+       "<title>54</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"608,-344 468,-344 468,-302 608,-302 608,-344\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"468,-302 468,-344 515,-344 515,-302 468,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"480\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"515,-323 515,-344 563,-344 563,-323 515,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"525\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"563,-323 563,-344 608,-344 608,-323 563,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"515,-302 515,-323 563,-323 563,-302 515,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"520\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"563,-302 563,-323 608,-323 608,-302 563,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 53&#45;&gt;54 -->\n",
+       "<g id=\"edge59\" class=\"edge\">\n",
+       "<title>53&#45;&gt;54</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-379.63C538,-371.82 538,-362.73 538,-354.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-354.16 538,-344.16 534.5,-354.16 541.5,-354.16\"/>\n",
+       "</g>\n",
+       "<!-- 55 -->\n",
+       "<g id=\"node56\" class=\"node\">\n",
+       "<title>55</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"609.5,-266 466.5,-266 466.5,-224 609.5,-224 609.5,-266\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"467,-224 467,-266 517,-266 517,-224 467,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"472\" y=\"-248\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"473.5\" y=\"-237\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"517,-245 517,-266 565,-266 565,-245 517,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"527\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"565,-245 565,-266 610,-266 610,-245 565,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"570\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"517,-224 517,-245 565,-245 565,-224 517,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"522\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"565,-224 565,-245 610,-245 610,-224 565,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"570\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 54&#45;&gt;55 -->\n",
+       "<g id=\"edge60\" class=\"edge\">\n",
+       "<title>54&#45;&gt;55</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-301.63C538,-293.82 538,-284.73 538,-276.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-276.16 538,-266.16 534.5,-276.16 541.5,-276.16\"/>\n",
+       "</g>\n",
+       "<!-- 56 -->\n",
+       "<g id=\"node57\" class=\"node\">\n",
+       "<title>56</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"608,-188 468,-188 468,-146 608,-146 608,-188\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"468,-146 468,-188 515,-188 515,-146 468,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"476\" y=\"-170\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-159\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"515,-167 515,-188 563,-188 563,-167 515,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"525\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"563,-167 563,-188 608,-188 608,-167 563,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"515,-146 515,-167 563,-167 563,-146 515,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"520\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"563,-146 563,-167 608,-167 608,-146 563,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 55&#45;&gt;56 -->\n",
+       "<g id=\"edge61\" class=\"edge\">\n",
+       "<title>55&#45;&gt;56</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-223.63C538,-215.82 538,-206.73 538,-198.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-198.16 538,-188.16 534.5,-198.16 541.5,-198.16\"/>\n",
+       "</g>\n",
+       "<!-- 57 -->\n",
+       "<g id=\"node58\" class=\"node\">\n",
+       "<title>57</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"617,-110 459,-110 459,-68 617,-68 617,-110\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"459,-68 459,-110 524,-110 524,-68 459,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"464\" y=\"-92\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-81\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"524,-89 524,-110 572,-110 572,-89 524,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"534\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"572,-89 572,-110 617,-110 617,-89 572,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"577\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"524,-68 524,-89 572,-89 572,-68 524,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"529\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"572,-68 572,-89 617,-89 617,-68 572,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"577\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 56&#45;&gt;57 -->\n",
+       "<g id=\"edge62\" class=\"edge\">\n",
+       "<title>56&#45;&gt;57</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-145.63C538,-137.82 538,-128.73 538,-120.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-120.16 538,-110.16 534.5,-120.16 541.5,-120.16\"/>\n",
+       "</g>\n",
+       "<!-- 58 -->\n",
+       "<g id=\"node59\" class=\"node\">\n",
+       "<title>58</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"597.5,-32 478.5,-32 478.5,0 597.5,0 597.5,-32\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"479,0 479,-32 556,-32 556,0 479,0\"/>\n",
+       "<text text-anchor=\"start\" x=\"484\" y=\"-19\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"499\" y=\"-8\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"556,0 556,-32 598,-32 598,0 556,0\"/>\n",
+       "<text text-anchor=\"start\" x=\"561\" y=\"-13.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64)</text>\n",
+       "</g>\n",
+       "<!-- 57&#45;&gt;58 -->\n",
+       "<g id=\"edge63\" class=\"edge\">\n",
+       "<title>57&#45;&gt;58</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M538,-67.84C538,-59.89 538,-50.66 538,-42.26\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"541.5,-42.24 538,-32.24 534.5,-42.24 541.5,-42.24\"/>\n",
+       "</g>\n",
+       "<!-- 68 -->\n",
+       "<g id=\"node69\" class=\"node\">\n",
+       "<title>68</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1071,-2216 931,-2216 931,-2174 1071,-2174 1071,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"931,-2174 931,-2216 978,-2216 978,-2174 931,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"939\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"936\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"978,-2195 978,-2216 1026,-2216 1026,-2195 978,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"988\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1026,-2195 1026,-2216 1071,-2216 1071,-2195 1026,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1031\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"978,-2174 978,-2195 1026,-2195 1026,-2174 978,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"983\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1026,-2174 1026,-2195 1071,-2195 1071,-2174 1026,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1031\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 59&#45;&gt;68 -->\n",
+       "<g id=\"edge71\" class=\"edge\">\n",
+       "<title>59&#45;&gt;68</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M936.58,-2251.63C946.82,-2242.79 958.96,-2232.3 969.93,-2222.84\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"972.38,-2225.34 977.66,-2216.16 967.8,-2220.05 972.38,-2225.34\"/>\n",
+       "</g>\n",
+       "<!-- 61 -->\n",
+       "<g id=\"node62\" class=\"node\">\n",
+       "<title>61</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1457.5,-2216 1282.5,-2216 1282.5,-2174 1457.5,-2174 1457.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1283,-2174 1283,-2216 1353,-2216 1353,-2174 1283,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1288\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1299.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2195 1353,-2216 1401,-2216 1401,-2195 1353,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2195 1401,-2216 1458,-2216 1458,-2195 1401,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2174 1353,-2195 1401,-2195 1401,-2174 1353,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2174 1401,-2195 1458,-2195 1458,-2174 1401,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "</g>\n",
+       "<!-- 60&#45;&gt;61 -->\n",
+       "<g id=\"edge64\" class=\"edge\">\n",
+       "<title>60&#45;&gt;61</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2251.63C1370,-2243.82 1370,-2234.73 1370,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2226.16 1370,-2216.16 1366.5,-2226.16 1373.5,-2226.16\"/>\n",
+       "</g>\n",
+       "<!-- 63 -->\n",
+       "<g id=\"node64\" class=\"node\">\n",
+       "<title>63</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1843.5,-2216 1668.5,-2216 1668.5,-2174 1843.5,-2174 1843.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1669,-2174 1669,-2216 1739,-2216 1739,-2174 1669,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1674\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1685.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1739,-2195 1739,-2216 1787,-2216 1787,-2195 1739,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1749\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1787,-2195 1787,-2216 1844,-2216 1844,-2195 1787,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1792\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1739,-2174 1739,-2195 1787,-2195 1787,-2174 1739,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1744\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1787,-2174 1787,-2195 1844,-2195 1844,-2174 1787,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1792\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "</g>\n",
+       "<!-- 60&#45;&gt;63 -->\n",
+       "<g id=\"edge65\" class=\"edge\">\n",
+       "<title>60&#45;&gt;63</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1457.63,-2255.3C1512.5,-2244.78 1584.93,-2230.77 1658.35,-2216.12\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1659.12,-2219.54 1668.24,-2214.14 1657.75,-2212.67 1659.12,-2219.54\"/>\n",
+       "</g>\n",
+       "<!-- 65 -->\n",
+       "<g id=\"node66\" class=\"node\">\n",
+       "<title>65</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1264.5,-2216 1089.5,-2216 1089.5,-2174 1264.5,-2174 1264.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1090,-2174 1090,-2216 1160,-2216 1160,-2174 1090,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1095\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1106.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1160,-2195 1160,-2216 1208,-2216 1208,-2195 1160,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1170\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1208,-2195 1208,-2216 1265,-2216 1265,-2195 1208,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1213\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1160,-2174 1160,-2195 1208,-2195 1208,-2174 1160,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1165\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1208,-2174 1208,-2195 1265,-2195 1265,-2174 1208,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1222\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 0) </text>\n",
+       "</g>\n",
+       "<!-- 60&#45;&gt;65 -->\n",
+       "<g id=\"edge66\" class=\"edge\">\n",
+       "<title>60&#45;&gt;65</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1319.04,-2251.93C1293.97,-2242.06 1263.59,-2230.1 1237.37,-2219.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1238.5,-2216.46 1227.91,-2216.05 1235.93,-2222.97 1238.5,-2216.46\"/>\n",
+       "</g>\n",
+       "<!-- 66 -->\n",
+       "<g id=\"node67\" class=\"node\">\n",
+       "<title>66</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1650.5,-2216 1475.5,-2216 1475.5,-2174 1650.5,-2174 1650.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1476,-2174 1476,-2216 1546,-2216 1546,-2174 1476,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1481\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1492.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1546,-2195 1546,-2216 1594,-2216 1594,-2195 1546,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1556\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1594,-2195 1594,-2216 1651,-2216 1651,-2195 1594,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1599\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1546,-2174 1546,-2195 1594,-2195 1594,-2174 1546,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1551\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1594,-2174 1594,-2195 1651,-2195 1651,-2174 1594,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1605\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "</g>\n",
+       "<!-- 60&#45;&gt;66 -->\n",
+       "<g id=\"edge67\" class=\"edge\">\n",
+       "<title>60&#45;&gt;66</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1420.96,-2251.93C1446.03,-2242.06 1476.41,-2230.1 1502.63,-2219.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1504.07,-2222.97 1512.09,-2216.05 1501.5,-2216.46 1504.07,-2222.97\"/>\n",
+       "</g>\n",
+       "<!-- 62 -->\n",
+       "<g id=\"node63\" class=\"node\">\n",
+       "<title>62</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1450.5,-2138 1289.5,-2138 1289.5,-2096 1450.5,-2096 1450.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1290,-2096 1290,-2138 1340,-2138 1340,-2096 1290,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1295\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1296.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2117 1340,-2138 1388,-2138 1388,-2117 1340,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1350\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2117 1388,-2138 1451,-2138 1451,-2117 1388,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1399\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2096 1340,-2117 1388,-2117 1388,-2096 1340,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1345\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2096 1388,-2117 1451,-2117 1451,-2096 1388,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "</g>\n",
+       "<!-- 61&#45;&gt;62 -->\n",
+       "<g id=\"edge68\" class=\"edge\">\n",
+       "<title>61&#45;&gt;62</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2173.63C1370,-2165.82 1370,-2156.73 1370,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2148.16 1370,-2138.16 1366.5,-2148.16 1373.5,-2148.16\"/>\n",
+       "</g>\n",
+       "<!-- 72 -->\n",
+       "<g id=\"node73\" class=\"node\">\n",
+       "<title>72</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-2060 1291,-2060 1291,-2018 1449,-2018 1449,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-2018 1291,-2060 1338,-2060 1338,-2018 1291,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1299\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2039 1338,-2060 1386,-2060 1386,-2039 1338,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2039 1386,-2060 1449,-2060 1449,-2039 1386,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2018 1338,-2039 1386,-2039 1386,-2018 1338,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1343\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2018 1386,-2039 1449,-2039 1449,-2018 1386,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 62&#45;&gt;72 -->\n",
+       "<g id=\"edge76\" class=\"edge\">\n",
+       "<title>62&#45;&gt;72</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2095.63C1370,-2087.82 1370,-2078.73 1370,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2070.16 1370,-2060.16 1366.5,-2070.16 1373.5,-2070.16\"/>\n",
+       "</g>\n",
+       "<!-- 64 -->\n",
+       "<g id=\"node65\" class=\"node\">\n",
+       "<title>64</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1838.5,-2138 1671.5,-2138 1671.5,-2096 1838.5,-2096 1838.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1672,-2096 1672,-2138 1722,-2138 1722,-2096 1672,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1677\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1678.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1722,-2117 1722,-2138 1770,-2138 1770,-2117 1722,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1732\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1770,-2117 1770,-2138 1839,-2138 1839,-2117 1770,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1781\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1722,-2096 1722,-2117 1770,-2117 1770,-2096 1722,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1727\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1770,-2096 1770,-2117 1839,-2117 1839,-2096 1770,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1775\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "</g>\n",
+       "<!-- 63&#45;&gt;64 -->\n",
+       "<g id=\"edge69\" class=\"edge\">\n",
+       "<title>63&#45;&gt;64</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1755.73,-2173.63C1755.63,-2165.82 1755.51,-2156.73 1755.4,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1758.9,-2148.11 1755.27,-2138.16 1751.9,-2148.2 1758.9,-2148.11\"/>\n",
+       "</g>\n",
+       "<!-- 84 -->\n",
+       "<g id=\"node85\" class=\"node\">\n",
+       "<title>84</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1834,-2060 1670,-2060 1670,-2018 1834,-2018 1834,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1670,-2018 1670,-2060 1717,-2060 1717,-2018 1670,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1678\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1675\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1717,-2039 1717,-2060 1765,-2060 1765,-2039 1717,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1727\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1765,-2039 1765,-2060 1834,-2060 1834,-2039 1765,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1770\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1717,-2018 1717,-2039 1765,-2039 1765,-2018 1717,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1722\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1765,-2018 1765,-2039 1834,-2039 1834,-2018 1765,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1770\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 64&#45;&gt;84 -->\n",
+       "<g id=\"edge90\" class=\"edge\">\n",
+       "<title>64&#45;&gt;84</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1754.2,-2095.63C1753.89,-2087.82 1753.53,-2078.73 1753.19,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1756.69,-2070.01 1752.8,-2060.16 1749.69,-2070.29 1756.69,-2070.01\"/>\n",
+       "</g>\n",
+       "<!-- 67 -->\n",
+       "<g id=\"node68\" class=\"node\">\n",
+       "<title>67</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1635.5,-2138 1486.5,-2138 1486.5,-2096 1635.5,-2096 1635.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1487,-2096 1487,-2138 1537,-2138 1537,-2096 1487,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1492\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1493.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1537,-2117 1537,-2138 1585,-2138 1585,-2117 1537,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1547\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1585,-2117 1585,-2138 1636,-2138 1636,-2117 1585,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1593\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1537,-2096 1537,-2117 1585,-2117 1585,-2096 1537,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1542\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1585,-2096 1585,-2117 1636,-2117 1636,-2096 1585,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1590\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 66&#45;&gt;67 -->\n",
+       "<g id=\"edge70\" class=\"edge\">\n",
+       "<title>66&#45;&gt;67</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1562.46,-2173.63C1562.26,-2165.82 1562.02,-2156.73 1561.79,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1565.29,-2148.06 1561.53,-2138.16 1558.29,-2148.25 1565.29,-2148.06\"/>\n",
+       "</g>\n",
+       "<!-- 80 -->\n",
+       "<g id=\"node81\" class=\"node\">\n",
+       "<title>80</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1633,-2060 1487,-2060 1487,-2018 1633,-2018 1633,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1487,-2018 1487,-2060 1534,-2060 1534,-2018 1487,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1495\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1492\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1534,-2039 1534,-2060 1582,-2060 1582,-2039 1534,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1544\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1582,-2039 1582,-2060 1633,-2060 1633,-2039 1582,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1587\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1534,-2018 1534,-2039 1582,-2039 1582,-2018 1534,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1539\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1582,-2018 1582,-2039 1633,-2039 1633,-2018 1582,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1587\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 67&#45;&gt;80 -->\n",
+       "<g id=\"edge85\" class=\"edge\">\n",
+       "<title>67&#45;&gt;80</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1560.73,-2095.63C1560.63,-2087.82 1560.51,-2078.73 1560.4,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1563.9,-2070.11 1560.27,-2060.16 1556.9,-2070.2 1563.9,-2070.11\"/>\n",
+       "</g>\n",
+       "<!-- 69 -->\n",
+       "<g id=\"node70\" class=\"node\">\n",
+       "<title>69</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1101.5,-2138 958.5,-2138 958.5,-2096 1101.5,-2096 1101.5,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"959,-2096 959,-2138 1009,-2138 1009,-2096 959,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"964\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"965.5\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1009,-2117 1009,-2138 1057,-2138 1057,-2117 1009,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1019\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1057,-2117 1057,-2138 1102,-2138 1102,-2117 1057,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1062\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1009,-2096 1009,-2117 1057,-2117 1057,-2096 1009,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1014\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1057,-2096 1057,-2117 1102,-2117 1102,-2096 1057,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1062\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 68&#45;&gt;69 -->\n",
+       "<g id=\"edge72\" class=\"edge\">\n",
+       "<title>68&#45;&gt;69</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1008.77,-2173.63C1011.82,-2165.65 1015.38,-2156.33 1018.7,-2147.62\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1022.01,-2148.75 1022.31,-2138.16 1015.47,-2146.25 1022.01,-2148.75\"/>\n",
+       "</g>\n",
+       "<!-- 70 -->\n",
+       "<g id=\"node71\" class=\"node\">\n",
+       "<title>70</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1138,-2060 980,-2060 980,-2018 1138,-2018 1138,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"980,-2018 980,-2060 1045,-2060 1045,-2018 980,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"985\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"994\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1045,-2039 1045,-2060 1093,-2060 1093,-2039 1045,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1055\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1093,-2039 1093,-2060 1138,-2060 1138,-2039 1093,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1098\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1045,-2018 1045,-2039 1093,-2039 1093,-2018 1045,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1050\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1093,-2018 1093,-2039 1138,-2039 1138,-2018 1093,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1098\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 69&#45;&gt;70 -->\n",
+       "<g id=\"edge73\" class=\"edge\">\n",
+       "<title>69&#45;&gt;70</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1037.77,-2095.63C1040.82,-2087.65 1044.38,-2078.33 1047.7,-2069.62\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1051.01,-2070.75 1051.31,-2060.16 1044.47,-2068.25 1051.01,-2070.75\"/>\n",
+       "</g>\n",
+       "<!-- 71 -->\n",
+       "<g id=\"node72\" class=\"node\">\n",
+       "<title>71</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1256,-1436 1116,-1436 1116,-1394 1256,-1394 1256,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1116,-1394 1116,-1436 1163,-1436 1163,-1394 1116,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1128\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1121\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1163,-1415 1163,-1436 1211,-1436 1211,-1415 1163,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1173\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1211,-1415 1211,-1436 1256,-1436 1256,-1415 1211,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1216\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1163,-1394 1163,-1415 1211,-1415 1211,-1394 1163,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1168\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1211,-1394 1211,-1415 1256,-1415 1256,-1394 1211,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1216\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 70&#45;&gt;71 -->\n",
+       "<g id=\"edge74\" class=\"edge\">\n",
+       "<title>70&#45;&gt;71</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1082.72,-2017.69C1111.73,-1990.47 1157,-1939.14 1157,-1884 1157,-1884 1157,-1884 1157,-1570 1157,-1526.25 1168.44,-1476.69 1176.95,-1445.89\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1180.34,-1446.77 1179.71,-1436.2 1173.61,-1444.86 1180.34,-1446.77\"/>\n",
+       "</g>\n",
+       "<!-- 101 -->\n",
+       "<g id=\"node102\" class=\"node\">\n",
+       "<title>101</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1584.5,-1202 1357.5,-1202 1357.5,-1160 1584.5,-1160 1584.5,-1202\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1358,-1160 1358,-1202 1405,-1202 1405,-1160 1358,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1373.5\" y=\"-1184\" font-family=\"Linux libertine\" font-size=\"10.00\">cat</text>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-1173\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1405,-1181 1405,-1202 1453,-1202 1453,-1181 1405,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"1415\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1453,-1181 1453,-1202 1585,-1202 1585,-1181 1453,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"1458\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10), 2 x (1, 32), (1, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1405,-1160 1405,-1181 1453,-1181 1453,-1160 1405,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1410\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1453,-1160 1453,-1181 1585,-1181 1585,-1160 1453,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501.5\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "</g>\n",
+       "<!-- 71&#45;&gt;101 -->\n",
+       "<g id=\"edge75\" class=\"edge\">\n",
+       "<title>71&#45;&gt;101</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1193.54,-1393.92C1208.22,-1357.49 1244.1,-1280.38 1299,-1238 1317.19,-1223.96 1339.03,-1213.37 1360.96,-1205.41\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1362.38,-1208.62 1370.68,-1202.04 1360.09,-1202 1362.38,-1208.62\"/>\n",
+       "</g>\n",
+       "<!-- 73 -->\n",
+       "<g id=\"node74\" class=\"node\">\n",
+       "<title>73</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1450.5,-1982 1289.5,-1982 1289.5,-1940 1450.5,-1940 1450.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1290,-1940 1290,-1982 1340,-1982 1340,-1940 1290,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1295\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1296.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-1961 1340,-1982 1388,-1982 1388,-1961 1340,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1350\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-1961 1388,-1982 1451,-1982 1451,-1961 1388,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-1940 1340,-1961 1388,-1961 1388,-1940 1340,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1345\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-1940 1388,-1961 1451,-1961 1451,-1940 1388,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 72&#45;&gt;73 -->\n",
+       "<g id=\"edge77\" class=\"edge\">\n",
+       "<title>72&#45;&gt;73</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2017.63C1370,-2009.82 1370,-2000.73 1370,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1992.16 1370,-1982.16 1366.5,-1992.16 1373.5,-1992.16\"/>\n",
+       "</g>\n",
+       "<!-- 74 -->\n",
+       "<g id=\"node75\" class=\"node\">\n",
+       "<title>74</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1458,-1904 1282,-1904 1282,-1862 1458,-1862 1458,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1282,-1862 1282,-1904 1347,-1904 1347,-1862 1282,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1287\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-1883 1347,-1904 1395,-1904 1395,-1883 1347,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-1883 1395,-1904 1458,-1904 1458,-1883 1395,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-1862 1347,-1883 1395,-1883 1395,-1862 1347,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1352\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-1862 1395,-1883 1458,-1883 1458,-1862 1395,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 73&#45;&gt;74 -->\n",
+       "<g id=\"edge78\" class=\"edge\">\n",
+       "<title>73&#45;&gt;74</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1939.63C1370,-1931.82 1370,-1922.73 1370,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1914.16 1370,-1904.16 1366.5,-1914.16 1373.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 75 -->\n",
+       "<g id=\"node76\" class=\"node\">\n",
+       "<title>75</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-1826 1291,-1826 1291,-1784 1449,-1784 1449,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-1784 1291,-1826 1338,-1826 1338,-1784 1291,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1303\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-1805 1338,-1826 1386,-1826 1386,-1805 1338,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-1805 1386,-1826 1449,-1826 1449,-1805 1386,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-1784 1338,-1805 1386,-1805 1386,-1784 1338,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1343\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-1784 1386,-1805 1449,-1805 1449,-1784 1386,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 74&#45;&gt;75 -->\n",
+       "<g id=\"edge79\" class=\"edge\">\n",
+       "<title>74&#45;&gt;75</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1861.63C1370,-1853.82 1370,-1844.73 1370,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1836.16 1370,-1826.16 1366.5,-1836.16 1373.5,-1836.16\"/>\n",
+       "</g>\n",
+       "<!-- 76 -->\n",
+       "<g id=\"node77\" class=\"node\">\n",
+       "<title>76</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1451,-1748 1293,-1748 1293,-1706 1451,-1706 1451,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1293,-1706 1293,-1748 1340,-1748 1340,-1706 1293,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1301\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1298\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-1727 1340,-1748 1388,-1748 1388,-1727 1340,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1350\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-1727 1388,-1748 1451,-1748 1451,-1727 1388,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-1706 1340,-1727 1388,-1727 1388,-1706 1340,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1345\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-1706 1388,-1727 1451,-1727 1451,-1706 1388,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 75&#45;&gt;76 -->\n",
+       "<g id=\"edge80\" class=\"edge\">\n",
+       "<title>75&#45;&gt;76</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370.54,-1783.63C1370.74,-1775.82 1370.98,-1766.73 1371.21,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1374.71,-1758.25 1371.47,-1748.16 1367.71,-1758.06 1374.71,-1758.25\"/>\n",
+       "</g>\n",
+       "<!-- 77 -->\n",
+       "<g id=\"node78\" class=\"node\">\n",
+       "<title>77</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1453.5,-1670 1292.5,-1670 1292.5,-1628 1453.5,-1628 1453.5,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1293,-1628 1293,-1670 1343,-1670 1343,-1628 1293,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1298\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1299.5\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1343,-1649 1343,-1670 1391,-1670 1391,-1649 1343,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1353\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1391,-1649 1391,-1670 1454,-1670 1454,-1649 1391,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1396\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1343,-1628 1343,-1649 1391,-1649 1391,-1628 1343,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1391,-1628 1391,-1649 1454,-1649 1454,-1628 1391,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1396\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 76&#45;&gt;77 -->\n",
+       "<g id=\"edge81\" class=\"edge\">\n",
+       "<title>76&#45;&gt;77</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1372.27,-1705.63C1372.37,-1697.82 1372.49,-1688.73 1372.6,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1376.1,-1680.2 1372.73,-1670.16 1369.1,-1680.11 1376.1,-1680.2\"/>\n",
+       "</g>\n",
+       "<!-- 78 -->\n",
+       "<g id=\"node79\" class=\"node\">\n",
+       "<title>78</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1470,-1592 1294,-1592 1294,-1550 1470,-1550 1470,-1592\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1294,-1550 1294,-1592 1359,-1592 1359,-1550 1294,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1299\" y=\"-1574\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1308\" y=\"-1563\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1359,-1571 1359,-1592 1407,-1592 1407,-1571 1359,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"1369\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1407,-1571 1407,-1592 1470,-1592 1470,-1571 1407,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"1412\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1359,-1550 1359,-1571 1407,-1571 1407,-1550 1359,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1364\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1407,-1550 1407,-1571 1470,-1571 1470,-1550 1407,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1412\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 77&#45;&gt;78 -->\n",
+       "<g id=\"edge82\" class=\"edge\">\n",
+       "<title>77&#45;&gt;78</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1375.41,-1627.63C1376.34,-1619.82 1377.41,-1610.73 1378.43,-1602.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1381.91,-1602.5 1379.61,-1592.16 1374.96,-1601.68 1381.91,-1602.5\"/>\n",
+       "</g>\n",
+       "<!-- 79 -->\n",
+       "<g id=\"node80\" class=\"node\">\n",
+       "<title>79</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1461,-1514 1303,-1514 1303,-1472 1461,-1472 1461,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1303,-1472 1303,-1514 1350,-1514 1350,-1472 1303,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1315\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1308\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1350,-1493 1350,-1514 1398,-1514 1398,-1493 1350,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"1360\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-1493 1398,-1514 1461,-1514 1461,-1493 1398,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1350,-1472 1350,-1493 1398,-1493 1398,-1472 1350,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1355\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-1472 1398,-1493 1461,-1493 1461,-1472 1398,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 78&#45;&gt;79 -->\n",
+       "<g id=\"edge83\" class=\"edge\">\n",
+       "<title>78&#45;&gt;79</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1382,-1549.63C1382,-1541.82 1382,-1532.73 1382,-1524.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1385.5,-1524.16 1382,-1514.16 1378.5,-1524.16 1385.5,-1524.16\"/>\n",
+       "</g>\n",
+       "<!-- 92 -->\n",
+       "<g id=\"node93\" class=\"node\">\n",
+       "<title>92</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1467.5,-1436 1304.5,-1436 1304.5,-1394 1467.5,-1394 1467.5,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1305,-1394 1305,-1436 1357,-1436 1357,-1394 1305,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1310\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"1312.5\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1357,-1415 1357,-1436 1405,-1436 1405,-1415 1357,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1367\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1405,-1415 1405,-1436 1468,-1436 1468,-1415 1405,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1410\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1357,-1394 1357,-1415 1405,-1415 1405,-1394 1357,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1362\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1405,-1394 1405,-1415 1468,-1415 1468,-1394 1405,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1410\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "</g>\n",
+       "<!-- 79&#45;&gt;92 -->\n",
+       "<g id=\"edge84\" class=\"edge\">\n",
+       "<title>79&#45;&gt;92</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1383.07,-1471.63C1383.48,-1463.82 1383.96,-1454.73 1384.41,-1446.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1387.91,-1446.33 1384.94,-1436.16 1380.92,-1445.96 1387.91,-1446.33\"/>\n",
+       "</g>\n",
+       "<!-- 81 -->\n",
+       "<g id=\"node82\" class=\"node\">\n",
+       "<title>81</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1633.5,-1982 1484.5,-1982 1484.5,-1940 1633.5,-1940 1633.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1485,-1940 1485,-1982 1535,-1982 1535,-1940 1485,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1490\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1491.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1535,-1961 1535,-1982 1583,-1982 1583,-1961 1535,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1545\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1583,-1961 1583,-1982 1634,-1982 1634,-1961 1583,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1588\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1535,-1940 1535,-1961 1583,-1961 1583,-1940 1535,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1540\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1583,-1940 1583,-1961 1634,-1961 1634,-1940 1583,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1588\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 80&#45;&gt;81 -->\n",
+       "<g id=\"edge86\" class=\"edge\">\n",
+       "<title>80&#45;&gt;81</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1559.73,-2017.63C1559.63,-2009.82 1559.51,-2000.73 1559.4,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1562.9,-1992.11 1559.27,-1982.16 1555.9,-1992.2 1562.9,-1992.11\"/>\n",
+       "</g>\n",
+       "<!-- 82 -->\n",
+       "<g id=\"node83\" class=\"node\">\n",
+       "<title>82</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1641,-1904 1477,-1904 1477,-1862 1641,-1862 1641,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1477,-1862 1477,-1904 1542,-1904 1542,-1862 1477,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1482\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1491\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1542,-1883 1542,-1904 1590,-1904 1590,-1883 1542,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1552\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1590,-1883 1590,-1904 1641,-1904 1641,-1883 1590,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1595\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1542,-1862 1542,-1883 1590,-1883 1590,-1862 1542,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1547\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1590,-1862 1590,-1883 1641,-1883 1641,-1862 1590,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1595\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 81&#45;&gt;82 -->\n",
+       "<g id=\"edge87\" class=\"edge\">\n",
+       "<title>81&#45;&gt;82</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1559,-1939.63C1559,-1931.82 1559,-1922.73 1559,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1562.5,-1914.16 1559,-1904.16 1555.5,-1914.16 1562.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 83 -->\n",
+       "<g id=\"node84\" class=\"node\">\n",
+       "<title>83</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1632,-1826 1486,-1826 1486,-1784 1632,-1784 1632,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1486,-1784 1486,-1826 1533,-1826 1533,-1784 1486,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1491\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1533,-1805 1533,-1826 1581,-1826 1581,-1805 1533,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1543\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1581,-1805 1581,-1826 1632,-1826 1632,-1805 1581,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1586\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1533,-1784 1533,-1805 1581,-1805 1581,-1784 1533,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1538\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1581,-1784 1581,-1805 1632,-1805 1632,-1784 1581,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1586\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 82&#45;&gt;83 -->\n",
+       "<g id=\"edge88\" class=\"edge\">\n",
+       "<title>82&#45;&gt;83</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1559,-1861.63C1559,-1853.82 1559,-1844.73 1559,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1562.5,-1836.16 1559,-1826.16 1555.5,-1836.16 1562.5,-1836.16\"/>\n",
+       "</g>\n",
+       "<!-- 95 -->\n",
+       "<g id=\"node96\" class=\"node\">\n",
+       "<title>95</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1633.5,-1748 1482.5,-1748 1482.5,-1706 1633.5,-1706 1633.5,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1483,-1706 1483,-1748 1535,-1748 1535,-1706 1483,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1488\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"1490.5\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1535,-1727 1535,-1748 1583,-1748 1583,-1727 1535,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1545\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1583,-1727 1583,-1748 1634,-1748 1634,-1727 1583,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1588\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1535,-1706 1535,-1727 1583,-1727 1583,-1706 1535,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1540\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1583,-1706 1583,-1727 1634,-1727 1634,-1706 1583,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1588\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "</g>\n",
+       "<!-- 83&#45;&gt;95 -->\n",
+       "<g id=\"edge89\" class=\"edge\">\n",
+       "<title>83&#45;&gt;95</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1558.73,-1783.63C1558.63,-1775.82 1558.51,-1766.73 1558.4,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1561.9,-1758.11 1558.27,-1748.16 1554.9,-1758.2 1561.9,-1758.11\"/>\n",
+       "</g>\n",
+       "<!-- 85 -->\n",
+       "<g id=\"node86\" class=\"node\">\n",
+       "<title>85</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1834.5,-1982 1667.5,-1982 1667.5,-1940 1834.5,-1940 1834.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1668,-1940 1668,-1982 1718,-1982 1718,-1940 1668,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1673\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1674.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1718,-1961 1718,-1982 1766,-1982 1766,-1961 1718,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1728\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1766,-1961 1766,-1982 1835,-1982 1835,-1961 1766,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1771\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1718,-1940 1718,-1961 1766,-1961 1766,-1940 1718,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1723\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1766,-1940 1766,-1961 1835,-1961 1835,-1940 1766,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1771\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 84&#45;&gt;85 -->\n",
+       "<g id=\"edge91\" class=\"edge\">\n",
+       "<title>84&#45;&gt;85</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1751.73,-2017.63C1751.63,-2009.82 1751.51,-2000.73 1751.4,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1754.9,-1992.11 1751.27,-1982.16 1747.9,-1992.2 1754.9,-1992.11\"/>\n",
+       "</g>\n",
+       "<!-- 86 -->\n",
+       "<g id=\"node87\" class=\"node\">\n",
+       "<title>86</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1842,-1904 1660,-1904 1660,-1862 1842,-1862 1842,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1660,-1862 1660,-1904 1725,-1904 1725,-1862 1660,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1665\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1674\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1725,-1883 1725,-1904 1773,-1904 1773,-1883 1725,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1735\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1773,-1883 1773,-1904 1842,-1904 1842,-1883 1773,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1778\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1725,-1862 1725,-1883 1773,-1883 1773,-1862 1725,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1730\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1773,-1862 1773,-1883 1842,-1883 1842,-1862 1773,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1778\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 85&#45;&gt;86 -->\n",
+       "<g id=\"edge92\" class=\"edge\">\n",
+       "<title>85&#45;&gt;86</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1751,-1939.63C1751,-1931.82 1751,-1922.73 1751,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1754.5,-1914.16 1751,-1904.16 1747.5,-1914.16 1754.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 87 -->\n",
+       "<g id=\"node88\" class=\"node\">\n",
+       "<title>87</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1832,-1826 1668,-1826 1668,-1784 1832,-1784 1832,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1668,-1784 1668,-1826 1715,-1826 1715,-1784 1668,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1680\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1673\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1715,-1805 1715,-1826 1763,-1826 1763,-1805 1715,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1725\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1763,-1805 1763,-1826 1832,-1826 1832,-1805 1763,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1768\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1715,-1784 1715,-1805 1763,-1805 1763,-1784 1715,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1720\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1763,-1784 1763,-1805 1832,-1805 1832,-1784 1763,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1768\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 86&#45;&gt;87 -->\n",
+       "<g id=\"edge93\" class=\"edge\">\n",
+       "<title>86&#45;&gt;87</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1750.73,-1861.63C1750.63,-1853.82 1750.51,-1844.73 1750.4,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1753.9,-1836.11 1750.27,-1826.16 1746.9,-1836.2 1753.9,-1836.11\"/>\n",
+       "</g>\n",
+       "<!-- 88 -->\n",
+       "<g id=\"node89\" class=\"node\">\n",
+       "<title>88</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1831,-1748 1667,-1748 1667,-1706 1831,-1706 1831,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1667,-1706 1667,-1748 1714,-1748 1714,-1706 1667,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1675\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1672\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1714,-1727 1714,-1748 1762,-1748 1762,-1727 1714,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1724\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1762,-1727 1762,-1748 1831,-1748 1831,-1727 1762,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1767\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1714,-1706 1714,-1727 1762,-1727 1762,-1706 1714,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1719\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1762,-1706 1762,-1727 1831,-1727 1831,-1706 1762,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1767\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 87&#45;&gt;88 -->\n",
+       "<g id=\"edge94\" class=\"edge\">\n",
+       "<title>87&#45;&gt;88</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1749.73,-1783.63C1749.63,-1775.82 1749.51,-1766.73 1749.4,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1752.9,-1758.11 1749.27,-1748.16 1745.9,-1758.2 1752.9,-1758.11\"/>\n",
+       "</g>\n",
+       "<!-- 89 -->\n",
+       "<g id=\"node90\" class=\"node\">\n",
+       "<title>89</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1831.5,-1670 1664.5,-1670 1664.5,-1628 1831.5,-1628 1831.5,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1665,-1628 1665,-1670 1715,-1670 1715,-1628 1665,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1670\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1671.5\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1715,-1649 1715,-1670 1763,-1670 1763,-1649 1715,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1725\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1763,-1649 1763,-1670 1832,-1670 1832,-1649 1763,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1768\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1715,-1628 1715,-1649 1763,-1649 1763,-1628 1715,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1720\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1763,-1628 1763,-1649 1832,-1649 1832,-1628 1763,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1768\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 88&#45;&gt;89 -->\n",
+       "<g id=\"edge95\" class=\"edge\">\n",
+       "<title>88&#45;&gt;89</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1748.73,-1705.63C1748.63,-1697.82 1748.51,-1688.73 1748.4,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1751.9,-1680.11 1748.27,-1670.16 1744.9,-1680.2 1751.9,-1680.11\"/>\n",
+       "</g>\n",
+       "<!-- 90 -->\n",
+       "<g id=\"node91\" class=\"node\">\n",
+       "<title>90</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1838,-1592 1656,-1592 1656,-1550 1838,-1550 1838,-1592\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1656,-1550 1656,-1592 1721,-1592 1721,-1550 1656,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1661\" y=\"-1574\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1670\" y=\"-1563\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1721,-1571 1721,-1592 1769,-1592 1769,-1571 1721,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"1731\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1769,-1571 1769,-1592 1838,-1592 1838,-1571 1769,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"1774\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1721,-1550 1721,-1571 1769,-1571 1769,-1550 1721,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1726\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1769,-1550 1769,-1571 1838,-1571 1838,-1550 1769,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1774\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 89&#45;&gt;90 -->\n",
+       "<g id=\"edge96\" class=\"edge\">\n",
+       "<title>89&#45;&gt;90</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1747.73,-1627.63C1747.63,-1619.82 1747.51,-1610.73 1747.4,-1602.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1750.9,-1602.11 1747.27,-1592.16 1743.9,-1602.2 1750.9,-1602.11\"/>\n",
+       "</g>\n",
+       "<!-- 91 -->\n",
+       "<g id=\"node92\" class=\"node\">\n",
+       "<title>91</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1819,-1514 1655,-1514 1655,-1472 1819,-1472 1819,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1655,-1472 1655,-1514 1702,-1514 1702,-1472 1655,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1667\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1660\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1702,-1493 1702,-1514 1750,-1514 1750,-1493 1702,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"1712\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1750,-1493 1750,-1514 1819,-1514 1819,-1493 1750,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"1755\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1702,-1472 1702,-1493 1750,-1493 1750,-1472 1702,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1707\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1750,-1472 1750,-1493 1819,-1493 1819,-1472 1750,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1755\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 90&#45;&gt;91 -->\n",
+       "<g id=\"edge97\" class=\"edge\">\n",
+       "<title>90&#45;&gt;91</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1744.32,-1549.63C1743.29,-1541.82 1742.1,-1532.73 1740.97,-1524.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1744.43,-1523.62 1739.65,-1514.16 1737.49,-1524.53 1744.43,-1523.62\"/>\n",
+       "</g>\n",
+       "<!-- 98 -->\n",
+       "<g id=\"node99\" class=\"node\">\n",
+       "<title>98</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1820.5,-1436 1651.5,-1436 1651.5,-1394 1820.5,-1394 1820.5,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1652,-1394 1652,-1436 1704,-1436 1704,-1394 1652,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1657\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"1659.5\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1704,-1415 1704,-1436 1752,-1436 1752,-1415 1704,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1714\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1752,-1415 1752,-1436 1821,-1436 1821,-1415 1752,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1757\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1704,-1394 1704,-1415 1752,-1415 1752,-1394 1704,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1709\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1752,-1394 1752,-1415 1821,-1415 1821,-1394 1752,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1757\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "</g>\n",
+       "<!-- 91&#45;&gt;98 -->\n",
+       "<g id=\"edge98\" class=\"edge\">\n",
+       "<title>91&#45;&gt;98</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1736.73,-1471.63C1736.63,-1463.82 1736.51,-1454.73 1736.4,-1446.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1739.9,-1446.11 1736.27,-1436.16 1732.9,-1446.2 1739.9,-1446.11\"/>\n",
+       "</g>\n",
+       "<!-- 93 -->\n",
+       "<g id=\"node94\" class=\"node\">\n",
+       "<title>93</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1477,-1358 1295,-1358 1295,-1316 1477,-1316 1477,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1295,-1316 1295,-1358 1366,-1358 1366,-1316 1295,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1300\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"1312\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1366,-1337 1366,-1358 1414,-1358 1414,-1337 1366,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1376\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1414,-1337 1414,-1358 1477,-1358 1477,-1337 1414,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1419\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1366,-1316 1366,-1337 1414,-1337 1414,-1316 1366,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1371\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1414,-1316 1414,-1337 1477,-1337 1477,-1316 1414,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1422\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 92&#45;&gt;93 -->\n",
+       "<g id=\"edge99\" class=\"edge\">\n",
+       "<title>92&#45;&gt;93</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1386,-1393.63C1386,-1385.82 1386,-1376.73 1386,-1368.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1389.5,-1368.16 1386,-1358.16 1382.5,-1368.16 1389.5,-1368.16\"/>\n",
+       "</g>\n",
+       "<!-- 94 -->\n",
+       "<g id=\"node95\" class=\"node\">\n",
+       "<title>94</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1464,-1280 1308,-1280 1308,-1238 1464,-1238 1464,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1308,-1238 1308,-1280 1359,-1280 1359,-1238 1308,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1313\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1315\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1359,-1259 1359,-1280 1407,-1280 1407,-1259 1359,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1369\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1407,-1259 1407,-1280 1464,-1280 1464,-1259 1407,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1412\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1359,-1238 1359,-1259 1407,-1259 1407,-1238 1359,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1364\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1407,-1238 1407,-1259 1464,-1259 1464,-1238 1407,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1418\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 93&#45;&gt;94 -->\n",
+       "<g id=\"edge100\" class=\"edge\">\n",
+       "<title>93&#45;&gt;94</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1386,-1315.63C1386,-1307.82 1386,-1298.73 1386,-1290.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1389.5,-1290.16 1386,-1280.16 1382.5,-1290.16 1389.5,-1290.16\"/>\n",
+       "</g>\n",
+       "<!-- 94&#45;&gt;101 -->\n",
+       "<g id=\"edge101\" class=\"edge\">\n",
+       "<title>94&#45;&gt;101</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1408.78,-1237.63C1418.67,-1228.79 1430.4,-1218.3 1440.99,-1208.84\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1443.33,-1211.43 1448.45,-1202.16 1438.67,-1206.22 1443.33,-1211.43\"/>\n",
+       "</g>\n",
+       "<!-- 96 -->\n",
+       "<g id=\"node97\" class=\"node\">\n",
+       "<title>96</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1643,-1670 1473,-1670 1473,-1628 1643,-1628 1643,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1473,-1628 1473,-1670 1544,-1670 1544,-1628 1473,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1478\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"1490\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1544,-1649 1544,-1670 1592,-1670 1592,-1649 1544,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1554\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1592,-1649 1592,-1670 1643,-1670 1643,-1649 1592,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1597\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1544,-1628 1544,-1649 1592,-1649 1592,-1628 1544,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1549\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1592,-1628 1592,-1649 1643,-1649 1643,-1628 1592,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1597\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "</g>\n",
+       "<!-- 95&#45;&gt;96 -->\n",
+       "<g id=\"edge102\" class=\"edge\">\n",
+       "<title>95&#45;&gt;96</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1558,-1705.63C1558,-1697.82 1558,-1688.73 1558,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1561.5,-1680.16 1558,-1670.16 1554.5,-1680.16 1561.5,-1680.16\"/>\n",
+       "</g>\n",
+       "<!-- 97 -->\n",
+       "<g id=\"node98\" class=\"node\">\n",
+       "<title>97</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1632,-1280 1482,-1280 1482,-1238 1632,-1238 1632,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1482,-1238 1482,-1280 1533,-1280 1533,-1238 1482,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1487\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1489\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1533,-1259 1533,-1280 1581,-1280 1581,-1259 1533,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1543\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1581,-1259 1581,-1280 1632,-1280 1632,-1259 1581,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1586\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1533,-1238 1533,-1259 1581,-1259 1581,-1238 1533,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1538\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1581,-1238 1581,-1259 1632,-1259 1632,-1238 1581,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1592\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3) </text>\n",
+       "</g>\n",
+       "<!-- 96&#45;&gt;97 -->\n",
+       "<g id=\"edge103\" class=\"edge\">\n",
+       "<title>96&#45;&gt;97</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1558,-1627.57C1558,-1598.08 1558,-1541.86 1558,-1494 1558,-1494 1558,-1494 1558,-1414 1558,-1371.19 1557.61,-1321.68 1557.32,-1290.61\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1560.81,-1290.39 1557.22,-1280.43 1553.81,-1290.46 1560.81,-1290.39\"/>\n",
+       "</g>\n",
+       "<!-- 97&#45;&gt;101 -->\n",
+       "<g id=\"edge104\" class=\"edge\">\n",
+       "<title>97&#45;&gt;101</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1533.95,-1237.63C1523.94,-1228.79 1512.08,-1218.3 1501.37,-1208.84\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1503.62,-1206.16 1493.81,-1202.16 1498.99,-1211.4 1503.62,-1206.16\"/>\n",
+       "</g>\n",
+       "<!-- 99 -->\n",
+       "<g id=\"node100\" class=\"node\">\n",
+       "<title>99</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1829,-1358 1641,-1358 1641,-1316 1829,-1316 1829,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1641,-1316 1641,-1358 1712,-1358 1712,-1316 1641,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1646\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"1658\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1712,-1337 1712,-1358 1760,-1358 1760,-1337 1712,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1722\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1760,-1337 1760,-1358 1829,-1358 1829,-1337 1760,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1765\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1712,-1316 1712,-1337 1760,-1337 1760,-1316 1712,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1717\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1760,-1316 1760,-1337 1829,-1337 1829,-1316 1760,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1771\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 98&#45;&gt;99 -->\n",
+       "<g id=\"edge105\" class=\"edge\">\n",
+       "<title>98&#45;&gt;99</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1735.73,-1393.63C1735.63,-1385.82 1735.51,-1376.73 1735.4,-1368.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1738.9,-1368.11 1735.27,-1358.16 1731.9,-1368.2 1738.9,-1368.11\"/>\n",
+       "</g>\n",
+       "<!-- 100 -->\n",
+       "<g id=\"node101\" class=\"node\">\n",
+       "<title>100</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1811,-1280 1655,-1280 1655,-1238 1811,-1238 1811,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1655,-1238 1655,-1280 1706,-1280 1706,-1238 1655,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1660\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1662\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1706,-1259 1706,-1280 1754,-1280 1754,-1259 1706,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1716\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1754,-1259 1754,-1280 1811,-1280 1811,-1259 1754,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1759\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1706,-1238 1706,-1259 1754,-1259 1754,-1238 1706,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1711\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1754,-1238 1754,-1259 1811,-1259 1811,-1238 1754,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1765\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 99&#45;&gt;100 -->\n",
+       "<g id=\"edge106\" class=\"edge\">\n",
+       "<title>99&#45;&gt;100</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1734.46,-1315.63C1734.26,-1307.82 1734.02,-1298.73 1733.79,-1290.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1737.29,-1290.06 1733.53,-1280.16 1730.29,-1290.25 1737.29,-1290.06\"/>\n",
+       "</g>\n",
+       "<!-- 100&#45;&gt;101 -->\n",
+       "<g id=\"edge107\" class=\"edge\">\n",
+       "<title>100&#45;&gt;101</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1663.82,-1237.93C1628.9,-1227.8 1586.37,-1215.47 1550.15,-1204.96\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1550.69,-1201.47 1540.11,-1202.05 1548.74,-1208.2 1550.69,-1201.47\"/>\n",
+       "</g>\n",
+       "<!-- 102 -->\n",
+       "<g id=\"node103\" class=\"node\">\n",
+       "<title>102</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1544,-1124 1398,-1124 1398,-1082 1544,-1082 1544,-1124\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-1082 1398,-1124 1445,-1124 1445,-1082 1398,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-1106\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-1095\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-1103 1445,-1124 1493,-1124 1493,-1103 1445,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"1455\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-1103 1493,-1124 1544,-1124 1544,-1103 1493,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-1082 1445,-1103 1493,-1103 1493,-1082 1445,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1450\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-1082 1493,-1103 1544,-1103 1544,-1082 1493,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 101&#45;&gt;102 -->\n",
+       "<g id=\"edge108\" class=\"edge\">\n",
+       "<title>101&#45;&gt;102</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-1159.63C1471,-1151.82 1471,-1142.73 1471,-1134.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-1134.16 1471,-1124.16 1467.5,-1134.16 1474.5,-1134.16\"/>\n",
+       "</g>\n",
+       "<!-- 103 -->\n",
+       "<g id=\"node104\" class=\"node\">\n",
+       "<title>103</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1553,-1046 1389,-1046 1389,-1004 1553,-1004 1553,-1046\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1389,-1004 1389,-1046 1454,-1046 1454,-1004 1389,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1394\" y=\"-1028\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-1017\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1454,-1025 1454,-1046 1502,-1046 1502,-1025 1454,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"1464\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1502,-1025 1502,-1046 1553,-1046 1553,-1025 1502,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"1507\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1454,-1004 1454,-1025 1502,-1025 1502,-1004 1454,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1459\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1502,-1004 1502,-1025 1553,-1025 1553,-1004 1502,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1507\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 102&#45;&gt;103 -->\n",
+       "<g id=\"edge109\" class=\"edge\">\n",
+       "<title>102&#45;&gt;103</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-1081.63C1471,-1073.82 1471,-1064.73 1471,-1056.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-1056.16 1471,-1046.16 1467.5,-1056.16 1474.5,-1056.16\"/>\n",
+       "</g>\n",
+       "<!-- 104 -->\n",
+       "<g id=\"node105\" class=\"node\">\n",
+       "<title>104</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1544,-968 1398,-968 1398,-926 1544,-926 1544,-968\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-926 1398,-968 1445,-968 1445,-926 1398,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1410\" y=\"-950\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-939\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-947 1445,-968 1493,-968 1493,-947 1445,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"1455\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-947 1493,-968 1544,-968 1544,-947 1493,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-926 1445,-947 1493,-947 1493,-926 1445,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1450\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-926 1493,-947 1544,-947 1544,-926 1493,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 103&#45;&gt;104 -->\n",
+       "<g id=\"edge110\" class=\"edge\">\n",
+       "<title>103&#45;&gt;104</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-1003.63C1471,-995.82 1471,-986.73 1471,-978.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-978.16 1471,-968.16 1467.5,-978.16 1474.5,-978.16\"/>\n",
+       "</g>\n",
+       "<!-- 105 -->\n",
+       "<g id=\"node106\" class=\"node\">\n",
+       "<title>105</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1545.5,-890 1396.5,-890 1396.5,-848 1545.5,-848 1545.5,-890\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1397,-848 1397,-890 1447,-890 1447,-848 1397,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1402\" y=\"-872\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1403.5\" y=\"-861\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1447,-869 1447,-890 1495,-890 1495,-869 1447,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"1457\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1495,-869 1495,-890 1546,-890 1546,-869 1495,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"1500\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1447,-848 1447,-869 1495,-869 1495,-848 1447,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1452\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1495,-848 1495,-869 1546,-869 1546,-848 1495,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1500\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 104&#45;&gt;105 -->\n",
+       "<g id=\"edge111\" class=\"edge\">\n",
+       "<title>104&#45;&gt;105</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-925.63C1471,-917.82 1471,-908.73 1471,-900.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-900.16 1471,-890.16 1467.5,-900.16 1474.5,-900.16\"/>\n",
+       "</g>\n",
+       "<!-- 106 -->\n",
+       "<g id=\"node107\" class=\"node\">\n",
+       "<title>106</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1544,-812 1398,-812 1398,-770 1544,-770 1544,-812\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-770 1398,-812 1445,-812 1445,-770 1398,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-794\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-783\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-791 1445,-812 1493,-812 1493,-791 1445,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"1455\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-791 1493,-812 1544,-812 1544,-791 1493,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-770 1445,-791 1493,-791 1493,-770 1445,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1450\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-770 1493,-791 1544,-791 1544,-770 1493,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 105&#45;&gt;106 -->\n",
+       "<g id=\"edge112\" class=\"edge\">\n",
+       "<title>105&#45;&gt;106</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-847.63C1471,-839.82 1471,-830.73 1471,-822.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-822.16 1471,-812.16 1467.5,-822.16 1474.5,-822.16\"/>\n",
+       "</g>\n",
+       "<!-- 107 -->\n",
+       "<g id=\"node108\" class=\"node\">\n",
+       "<title>107</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1553,-734 1389,-734 1389,-692 1553,-692 1553,-734\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1389,-692 1389,-734 1454,-734 1454,-692 1389,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1394\" y=\"-716\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-705\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1454,-713 1454,-734 1502,-734 1502,-713 1454,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"1464\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1502,-713 1502,-734 1553,-734 1553,-713 1502,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"1507\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1454,-692 1454,-713 1502,-713 1502,-692 1454,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1459\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1502,-692 1502,-713 1553,-713 1553,-692 1502,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1507\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 106&#45;&gt;107 -->\n",
+       "<g id=\"edge113\" class=\"edge\">\n",
+       "<title>106&#45;&gt;107</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-769.63C1471,-761.82 1471,-752.73 1471,-744.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-744.16 1471,-734.16 1467.5,-744.16 1474.5,-744.16\"/>\n",
+       "</g>\n",
+       "<!-- 108 -->\n",
+       "<g id=\"node109\" class=\"node\">\n",
+       "<title>108</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1544,-656 1398,-656 1398,-614 1544,-614 1544,-656\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-614 1398,-656 1445,-656 1445,-614 1398,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1410\" y=\"-638\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-627\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-635 1445,-656 1493,-656 1493,-635 1445,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1455\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-635 1493,-656 1544,-656 1544,-635 1493,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-614 1445,-635 1493,-635 1493,-614 1445,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1450\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-614 1493,-635 1544,-635 1544,-614 1493,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 107&#45;&gt;108 -->\n",
+       "<g id=\"edge114\" class=\"edge\">\n",
+       "<title>107&#45;&gt;108</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-691.63C1471,-683.82 1471,-674.73 1471,-666.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-666.16 1471,-656.16 1467.5,-666.16 1474.5,-666.16\"/>\n",
+       "</g>\n",
+       "<!-- 109 -->\n",
+       "<g id=\"node110\" class=\"node\">\n",
+       "<title>109</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1545.5,-578 1396.5,-578 1396.5,-536 1545.5,-536 1545.5,-578\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1397,-536 1397,-578 1447,-578 1447,-536 1397,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1402\" y=\"-560\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1403.5\" y=\"-549\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1447,-557 1447,-578 1495,-578 1495,-557 1447,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1457\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1495,-557 1495,-578 1546,-578 1546,-557 1495,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1500\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1447,-536 1447,-557 1495,-557 1495,-536 1447,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1452\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1495,-536 1495,-557 1546,-557 1546,-536 1495,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1500\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 108&#45;&gt;109 -->\n",
+       "<g id=\"edge115\" class=\"edge\">\n",
+       "<title>108&#45;&gt;109</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-613.63C1471,-605.82 1471,-596.73 1471,-588.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-588.16 1471,-578.16 1467.5,-588.16 1474.5,-588.16\"/>\n",
+       "</g>\n",
+       "<!-- 110 -->\n",
+       "<g id=\"node111\" class=\"node\">\n",
+       "<title>110</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1544,-500 1398,-500 1398,-458 1544,-458 1544,-500\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-458 1398,-500 1445,-500 1445,-458 1398,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-479 1445,-500 1493,-500 1493,-479 1445,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1455\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-479 1493,-500 1544,-500 1544,-479 1493,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-458 1445,-479 1493,-479 1493,-458 1445,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1450\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-458 1493,-479 1544,-479 1544,-458 1493,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 109&#45;&gt;110 -->\n",
+       "<g id=\"edge116\" class=\"edge\">\n",
+       "<title>109&#45;&gt;110</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-535.63C1471,-527.82 1471,-518.73 1471,-510.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-510.16 1471,-500.16 1467.5,-510.16 1474.5,-510.16\"/>\n",
+       "</g>\n",
+       "<!-- 111 -->\n",
+       "<g id=\"node112\" class=\"node\">\n",
+       "<title>111</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1550,-422 1392,-422 1392,-380 1550,-380 1550,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-380 1392,-422 1457,-422 1457,-380 1392,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1457,-401 1457,-422 1505,-422 1505,-401 1457,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1467\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1505,-401 1505,-422 1550,-422 1550,-401 1505,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1510\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1457,-380 1457,-401 1505,-401 1505,-380 1457,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1462\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1505,-380 1505,-401 1550,-401 1550,-380 1505,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1510\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 110&#45;&gt;111 -->\n",
+       "<g id=\"edge117\" class=\"edge\">\n",
+       "<title>110&#45;&gt;111</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-457.63C1471,-449.82 1471,-440.73 1471,-432.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-432.16 1471,-422.16 1467.5,-432.16 1474.5,-432.16\"/>\n",
+       "</g>\n",
+       "<!-- 112 -->\n",
+       "<g id=\"node113\" class=\"node\">\n",
+       "<title>112</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1541,-344 1401,-344 1401,-302 1541,-302 1541,-344\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-302 1401,-344 1448,-344 1448,-302 1401,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1413\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1448,-323 1448,-344 1496,-344 1496,-323 1448,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1458\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1496,-323 1496,-344 1541,-344 1541,-323 1496,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1448,-302 1448,-323 1496,-323 1496,-302 1448,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1453\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1496,-302 1496,-323 1541,-323 1541,-302 1496,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 111&#45;&gt;112 -->\n",
+       "<g id=\"edge118\" class=\"edge\">\n",
+       "<title>111&#45;&gt;112</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-379.63C1471,-371.82 1471,-362.73 1471,-354.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-354.16 1471,-344.16 1467.5,-354.16 1474.5,-354.16\"/>\n",
+       "</g>\n",
+       "<!-- 113 -->\n",
+       "<g id=\"node114\" class=\"node\">\n",
+       "<title>113</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1542.5,-266 1399.5,-266 1399.5,-224 1542.5,-224 1542.5,-266\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1400,-224 1400,-266 1450,-266 1450,-224 1400,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1405\" y=\"-248\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1406.5\" y=\"-237\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1450,-245 1450,-266 1498,-266 1498,-245 1450,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1460\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1498,-245 1498,-266 1543,-266 1543,-245 1498,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1503\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1450,-224 1450,-245 1498,-245 1498,-224 1450,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1455\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1498,-224 1498,-245 1543,-245 1543,-224 1498,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1503\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 112&#45;&gt;113 -->\n",
+       "<g id=\"edge119\" class=\"edge\">\n",
+       "<title>112&#45;&gt;113</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-301.63C1471,-293.82 1471,-284.73 1471,-276.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-276.16 1471,-266.16 1467.5,-276.16 1474.5,-276.16\"/>\n",
+       "</g>\n",
+       "<!-- 114 -->\n",
+       "<g id=\"node115\" class=\"node\">\n",
+       "<title>114</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1541,-188 1401,-188 1401,-146 1541,-146 1541,-188\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-146 1401,-188 1448,-188 1448,-146 1401,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-170\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-159\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1448,-167 1448,-188 1496,-188 1496,-167 1448,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1458\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1496,-167 1496,-188 1541,-188 1541,-167 1496,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1448,-146 1448,-167 1496,-167 1496,-146 1448,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1453\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1496,-146 1496,-167 1541,-167 1541,-146 1496,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1501\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 113&#45;&gt;114 -->\n",
+       "<g id=\"edge120\" class=\"edge\">\n",
+       "<title>113&#45;&gt;114</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-223.63C1471,-215.82 1471,-206.73 1471,-198.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-198.16 1471,-188.16 1467.5,-198.16 1474.5,-198.16\"/>\n",
+       "</g>\n",
+       "<!-- 115 -->\n",
+       "<g id=\"node116\" class=\"node\">\n",
+       "<title>115</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1550,-110 1392,-110 1392,-68 1550,-68 1550,-110\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-68 1392,-110 1457,-110 1457,-68 1392,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-92\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-81\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1457,-89 1457,-110 1505,-110 1505,-89 1457,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"1467\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1505,-89 1505,-110 1550,-110 1550,-89 1505,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"1510\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1457,-68 1457,-89 1505,-89 1505,-68 1457,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"1462\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1505,-68 1505,-89 1550,-89 1550,-68 1505,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"1510\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 114&#45;&gt;115 -->\n",
+       "<g id=\"edge121\" class=\"edge\">\n",
+       "<title>114&#45;&gt;115</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-145.63C1471,-137.82 1471,-128.73 1471,-120.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-120.16 1471,-110.16 1467.5,-120.16 1474.5,-120.16\"/>\n",
+       "</g>\n",
+       "<!-- 116 -->\n",
+       "<g id=\"node117\" class=\"node\">\n",
+       "<title>116</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"1530.5,-32 1411.5,-32 1411.5,0 1530.5,0 1530.5,-32\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1412,0 1412,-32 1489,-32 1489,0 1412,0\"/>\n",
+       "<text text-anchor=\"start\" x=\"1417\" y=\"-19\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"1432\" y=\"-8\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1489,0 1489,-32 1531,-32 1531,0 1489,0\"/>\n",
+       "<text text-anchor=\"start\" x=\"1494\" y=\"-13.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64)</text>\n",
+       "</g>\n",
+       "<!-- 115&#45;&gt;116 -->\n",
+       "<g id=\"edge122\" class=\"edge\">\n",
+       "<title>115&#45;&gt;116</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1471,-67.84C1471,-59.89 1471,-50.66 1471,-42.26\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1474.5,-42.24 1471,-32.24 1467.5,-42.24 1474.5,-42.24\"/>\n",
+       "</g>\n",
+       "</g>\n",
+       "</svg>\n"
+      ],
+      "text/plain": [
+       "<graphviz.graphs.Digraph at 0x7fc778e9f310>"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_graph = draw_graph(model.policy.mlp_extractor, input_data=torch.Tensor(y).unsqueeze(0).cuda(), device='cuda')\n",
+    "model_graph.visual_graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_graph = draw_graph(model.policy, input_data=torch.Tensor(y).unsqueeze(0).cuda(), device='cuda')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/svg+xml": [
+       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
+       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
+       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
+       "<!-- Generated by graphviz version 2.43.0 (0)\n",
+       " -->\n",
+       "<!-- Title: model Pages: 1 -->\n",
+       "<svg width=\"1969pt\" height=\"3072pt\"\n",
+       " viewBox=\"0.00 0.00 1968.50 3072.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
+       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 3068)\">\n",
+       "<title>model</title>\n",
+       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-3068 1964.5,-3068 1964.5,4 -4,4\"/>\n",
+       "<!-- 0 -->\n",
+       "<g id=\"node1\" class=\"node\">\n",
+       "<title>0</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"878,-3064 754,-3064 754,-3032 878,-3032 878,-3064\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"754,-3032 754,-3064 824,-3064 824,-3032 754,-3032\"/>\n",
+       "<text text-anchor=\"start\" x=\"759\" y=\"-3051\" font-family=\"Linux libertine\" font-size=\"10.00\">input&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"770.5\" y=\"-3040\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"824,-3032 824,-3064 878,-3064 878,-3032 824,-3032\"/>\n",
+       "<text text-anchor=\"start\" x=\"829\" y=\"-3045.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730)</text>\n",
+       "</g>\n",
+       "<!-- 1 -->\n",
+       "<g id=\"node2\" class=\"node\">\n",
+       "<title>1</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"892,-2996 740,-2996 740,-2954 892,-2954 892,-2996\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"740,-2954 740,-2996 787,-2996 787,-2954 740,-2954\"/>\n",
+       "<text text-anchor=\"start\" x=\"752.5\" y=\"-2978\" font-family=\"Linux libertine\" font-size=\"10.00\">float</text>\n",
+       "<text text-anchor=\"start\" x=\"745\" y=\"-2967\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-2975 787,-2996 835,-2996 835,-2975 787,-2975\"/>\n",
+       "<text text-anchor=\"start\" x=\"797\" y=\"-2983\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-2975 835,-2996 892,-2996 892,-2975 835,-2975\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-2983\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-2954 787,-2975 835,-2975 835,-2954 787,-2954\"/>\n",
+       "<text text-anchor=\"start\" x=\"792\" y=\"-2962\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-2954 835,-2975 892,-2975 892,-2954 835,-2954\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-2962\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "</g>\n",
+       "<!-- 0&#45;&gt;1 -->\n",
+       "<g id=\"edge1\" class=\"edge\">\n",
+       "<title>0&#45;&gt;1</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M816,-3031.94C816,-3024.45 816,-3015.12 816,-3006.24\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"819.5,-3006.16 816,-2996.16 812.5,-3006.16 819.5,-3006.16\"/>\n",
+       "</g>\n",
+       "<!-- 2 -->\n",
+       "<g id=\"node3\" class=\"node\">\n",
+       "<title>2</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"892,-2918 740,-2918 740,-2876 892,-2876 892,-2918\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"740,-2876 740,-2918 787,-2918 787,-2876 740,-2876\"/>\n",
+       "<text text-anchor=\"start\" x=\"746\" y=\"-2900\" font-family=\"Linux libertine\" font-size=\"10.00\">Flatten</text>\n",
+       "<text text-anchor=\"start\" x=\"745\" y=\"-2889\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-2897 787,-2918 835,-2918 835,-2897 787,-2897\"/>\n",
+       "<text text-anchor=\"start\" x=\"797\" y=\"-2905\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-2897 835,-2918 892,-2918 892,-2897 835,-2897\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-2905\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-2876 787,-2897 835,-2897 835,-2876 787,-2876\"/>\n",
+       "<text text-anchor=\"start\" x=\"792\" y=\"-2884\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-2876 835,-2897 892,-2897 892,-2876 835,-2876\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-2884\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "</g>\n",
+       "<!-- 1&#45;&gt;2 -->\n",
+       "<g id=\"edge2\" class=\"edge\">\n",
+       "<title>1&#45;&gt;2</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M816,-2953.63C816,-2945.82 816,-2936.73 816,-2928.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"819.5,-2928.16 816,-2918.16 812.5,-2928.16 819.5,-2928.16\"/>\n",
+       "</g>\n",
+       "<!-- 3 -->\n",
+       "<g id=\"node4\" class=\"node\">\n",
+       "<title>3</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"342.5,-2840 167.5,-2840 167.5,-2798 342.5,-2798 342.5,-2840\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"168,-2798 168,-2840 238,-2840 238,-2798 168,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"173\" y=\"-2822\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"184.5\" y=\"-2811\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"238,-2819 238,-2840 286,-2840 286,-2819 238,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"248\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"286,-2819 286,-2840 343,-2840 343,-2819 286,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"291\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"238,-2798 238,-2819 286,-2819 286,-2798 238,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"243\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"286,-2798 286,-2819 343,-2819 343,-2798 286,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"297\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;3 -->\n",
+       "<g id=\"edge3\" class=\"edge\">\n",
+       "<title>2&#45;&gt;3</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M739.76,-2885.67C639.38,-2872.07 462.8,-2848.15 352.72,-2833.24\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"352.99,-2829.74 342.61,-2831.87 352.05,-2836.68 352.99,-2829.74\"/>\n",
+       "</g>\n",
+       "<!-- 4 -->\n",
+       "<g id=\"node5\" class=\"node\">\n",
+       "<title>4</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"719.5,-2840 544.5,-2840 544.5,-2798 719.5,-2798 719.5,-2840\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"545,-2798 545,-2840 615,-2840 615,-2798 545,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"550\" y=\"-2822\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"561.5\" y=\"-2811\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2819 615,-2840 663,-2840 663,-2819 615,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"625\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2819 663,-2840 720,-2840 720,-2819 663,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2798 615,-2819 663,-2819 663,-2798 615,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"620\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2798 663,-2819 720,-2819 720,-2798 663,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;4 -->\n",
+       "<g id=\"edge4\" class=\"edge\">\n",
+       "<title>2&#45;&gt;4</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M767.17,-2875.83C743.48,-2866.05 714.84,-2854.22 690.04,-2843.97\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"691.31,-2840.71 680.73,-2840.13 688.64,-2847.18 691.31,-2840.71\"/>\n",
+       "</g>\n",
+       "<!-- 60 -->\n",
+       "<g id=\"node61\" class=\"node\">\n",
+       "<title>60</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1000.5,-2840 825.5,-2840 825.5,-2798 1000.5,-2798 1000.5,-2840\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"826,-2798 826,-2840 896,-2840 896,-2798 826,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"831\" y=\"-2822\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"842.5\" y=\"-2811\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"896,-2819 896,-2840 944,-2840 944,-2819 896,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"906\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-2819 944,-2840 1001,-2840 1001,-2819 944,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"949\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"896,-2798 896,-2819 944,-2819 944,-2798 896,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"901\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-2798 944,-2819 1001,-2819 1001,-2798 944,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"955\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;60 -->\n",
+       "<g id=\"edge5\" class=\"edge\">\n",
+       "<title>2&#45;&gt;60</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M841.74,-2875.83C853.14,-2866.9 866.71,-2856.27 878.94,-2846.69\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"881.43,-2849.18 887.14,-2840.26 877.11,-2843.67 881.43,-2849.18\"/>\n",
+       "</g>\n",
+       "<!-- 61 -->\n",
+       "<g id=\"node62\" class=\"node\">\n",
+       "<title>61</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1457.5,-2840 1282.5,-2840 1282.5,-2798 1457.5,-2798 1457.5,-2840\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1283,-2798 1283,-2840 1353,-2840 1353,-2798 1283,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"1288\" y=\"-2822\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1299.5\" y=\"-2811\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2819 1353,-2840 1401,-2840 1401,-2819 1353,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2819 1401,-2840 1458,-2840 1458,-2819 1401,-2819\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-2827\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6730) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2798 1353,-2819 1401,-2819 1401,-2798 1353,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2798 1401,-2819 1458,-2819 1458,-2798 1401,-2798\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-2806\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "</g>\n",
+       "<!-- 2&#45;&gt;61 -->\n",
+       "<g id=\"edge6\" class=\"edge\">\n",
+       "<title>2&#45;&gt;61</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M892.16,-2885.55C991.1,-2871.98 1163.95,-2848.27 1272.41,-2833.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1272.94,-2836.85 1282.37,-2832.02 1271.99,-2829.91 1272.94,-2836.85\"/>\n",
+       "</g>\n",
+       "<!-- 12 -->\n",
+       "<g id=\"node13\" class=\"node\">\n",
+       "<title>12</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"140,-2762 0,-2762 0,-2720 140,-2720 140,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"0,-2720 0,-2762 47,-2762 47,-2720 0,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"8\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"47,-2741 47,-2762 95,-2762 95,-2741 47,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"57\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"95,-2741 95,-2762 140,-2762 140,-2741 95,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"100\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"47,-2720 47,-2741 95,-2741 95,-2720 47,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"52\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"95,-2720 95,-2741 140,-2741 140,-2720 95,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"100\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 3&#45;&gt;12 -->\n",
+       "<g id=\"edge14\" class=\"edge\">\n",
+       "<title>3&#45;&gt;12</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M205.91,-2797.83C182.08,-2788.05 153.29,-2776.22 128.35,-2765.97\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"129.58,-2762.69 119,-2762.13 126.92,-2769.17 129.58,-2762.69\"/>\n",
+       "</g>\n",
+       "<!-- 5 -->\n",
+       "<g id=\"node6\" class=\"node\">\n",
+       "<title>5</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"526.5,-2762 351.5,-2762 351.5,-2720 526.5,-2720 526.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"352,-2720 352,-2762 422,-2762 422,-2720 352,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"357\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"368.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"422,-2741 422,-2762 470,-2762 470,-2741 422,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"432\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"470,-2741 470,-2762 527,-2762 527,-2741 470,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"475\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"422,-2720 422,-2741 470,-2741 470,-2720 422,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"427\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"470,-2720 470,-2741 527,-2741 527,-2720 470,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"478\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "</g>\n",
+       "<!-- 4&#45;&gt;5 -->\n",
+       "<g id=\"edge7\" class=\"edge\">\n",
+       "<title>4&#45;&gt;5</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M581.04,-2797.93C555.97,-2788.06 525.59,-2776.1 499.37,-2765.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"500.5,-2762.46 489.91,-2762.05 497.93,-2768.97 500.5,-2762.46\"/>\n",
+       "</g>\n",
+       "<!-- 7 -->\n",
+       "<g id=\"node8\" class=\"node\">\n",
+       "<title>7</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"912.5,-2762 737.5,-2762 737.5,-2720 912.5,-2720 912.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"738,-2720 738,-2762 808,-2762 808,-2720 738,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"743\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"754.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"808,-2741 808,-2762 856,-2762 856,-2741 808,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"818\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"856,-2741 856,-2762 913,-2762 913,-2741 856,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"861\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"808,-2720 808,-2741 856,-2741 856,-2720 808,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"813\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"856,-2720 856,-2741 913,-2741 913,-2720 856,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"861\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "</g>\n",
+       "<!-- 4&#45;&gt;7 -->\n",
+       "<g id=\"edge8\" class=\"edge\">\n",
+       "<title>4&#45;&gt;7</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M682.96,-2797.93C708.03,-2788.06 738.41,-2776.1 764.63,-2765.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"766.07,-2768.97 774.09,-2762.05 763.5,-2762.46 766.07,-2768.97\"/>\n",
+       "</g>\n",
+       "<!-- 9 -->\n",
+       "<g id=\"node10\" class=\"node\">\n",
+       "<title>9</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"333.5,-2762 158.5,-2762 158.5,-2720 333.5,-2720 333.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"159,-2720 159,-2762 229,-2762 229,-2720 159,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"164\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"175.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"229,-2741 229,-2762 277,-2762 277,-2741 229,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"239\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"277,-2741 277,-2762 334,-2762 334,-2741 277,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"282\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"229,-2720 229,-2741 277,-2741 277,-2720 229,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"234\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"277,-2720 277,-2741 334,-2741 334,-2720 277,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"291\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 0) </text>\n",
+       "</g>\n",
+       "<!-- 4&#45;&gt;9 -->\n",
+       "<g id=\"edge9\" class=\"edge\">\n",
+       "<title>4&#45;&gt;9</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M544.37,-2801.42C489.28,-2790.92 416.45,-2776.89 343.45,-2762.28\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"344.12,-2758.84 333.62,-2760.31 342.74,-2765.71 344.12,-2758.84\"/>\n",
+       "</g>\n",
+       "<!-- 10 -->\n",
+       "<g id=\"node11\" class=\"node\">\n",
+       "<title>10</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"719.5,-2762 544.5,-2762 544.5,-2720 719.5,-2720 719.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"545,-2720 545,-2762 615,-2762 615,-2720 545,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"550\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"561.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2741 615,-2762 663,-2762 663,-2741 615,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"625\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2741 663,-2762 720,-2762 720,-2741 663,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2720 615,-2741 663,-2741 663,-2720 615,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"620\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2720 663,-2741 720,-2741 720,-2720 663,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"674\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "</g>\n",
+       "<!-- 4&#45;&gt;10 -->\n",
+       "<g id=\"edge10\" class=\"edge\">\n",
+       "<title>4&#45;&gt;10</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2797.63C632,-2789.82 632,-2780.73 632,-2772.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2772.16 632,-2762.16 628.5,-2772.16 635.5,-2772.16\"/>\n",
+       "</g>\n",
+       "<!-- 6 -->\n",
+       "<g id=\"node7\" class=\"node\">\n",
+       "<title>6</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"519.5,-2684 358.5,-2684 358.5,-2642 519.5,-2642 519.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"359,-2642 359,-2684 409,-2684 409,-2642 359,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"364\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"365.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-2663 409,-2684 457,-2684 457,-2663 409,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"419\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"457,-2663 457,-2684 520,-2684 520,-2663 457,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"468\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-2642 409,-2663 457,-2663 457,-2642 409,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"414\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"457,-2642 457,-2663 520,-2663 520,-2642 457,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"462\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "</g>\n",
+       "<!-- 5&#45;&gt;6 -->\n",
+       "<g id=\"edge11\" class=\"edge\">\n",
+       "<title>5&#45;&gt;6</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M439,-2719.63C439,-2711.82 439,-2702.73 439,-2694.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"442.5,-2694.16 439,-2684.16 435.5,-2694.16 442.5,-2694.16\"/>\n",
+       "</g>\n",
+       "<!-- 16 -->\n",
+       "<g id=\"node17\" class=\"node\">\n",
+       "<title>16</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"520,-2606 362,-2606 362,-2564 520,-2564 520,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"362,-2564 362,-2606 409,-2606 409,-2564 362,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"370\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"367\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-2585 409,-2606 457,-2606 457,-2585 409,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"419\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"457,-2585 457,-2606 520,-2606 520,-2585 457,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"462\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-2564 409,-2585 457,-2585 457,-2564 409,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"414\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"457,-2564 457,-2585 520,-2585 520,-2564 457,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"462\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 6&#45;&gt;16 -->\n",
+       "<g id=\"edge19\" class=\"edge\">\n",
+       "<title>6&#45;&gt;16</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M439.54,-2641.63C439.74,-2633.82 439.98,-2624.73 440.21,-2616.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"443.71,-2616.25 440.47,-2606.16 436.71,-2616.06 443.71,-2616.25\"/>\n",
+       "</g>\n",
+       "<!-- 8 -->\n",
+       "<g id=\"node9\" class=\"node\">\n",
+       "<title>8</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"907.5,-2684 740.5,-2684 740.5,-2642 907.5,-2642 907.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"741,-2642 741,-2684 791,-2684 791,-2642 741,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"746\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"747.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-2663 791,-2684 839,-2684 839,-2663 791,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"801\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-2663 839,-2684 908,-2684 908,-2663 839,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"850\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-2642 791,-2663 839,-2663 839,-2642 791,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"796\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-2642 839,-2663 908,-2663 908,-2642 839,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "</g>\n",
+       "<!-- 7&#45;&gt;8 -->\n",
+       "<g id=\"edge12\" class=\"edge\">\n",
+       "<title>7&#45;&gt;8</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M824.73,-2719.63C824.63,-2711.82 824.51,-2702.73 824.4,-2694.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"827.9,-2694.11 824.27,-2684.16 820.9,-2694.2 827.9,-2694.11\"/>\n",
+       "</g>\n",
+       "<!-- 28 -->\n",
+       "<g id=\"node29\" class=\"node\">\n",
+       "<title>28</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"906,-2606 742,-2606 742,-2564 906,-2564 906,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"742,-2564 742,-2606 789,-2606 789,-2564 742,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"750\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"747\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"789,-2585 789,-2606 837,-2606 837,-2585 789,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"799\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"837,-2585 837,-2606 906,-2606 906,-2585 837,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"842\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"789,-2564 789,-2585 837,-2585 837,-2564 789,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"794\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"837,-2564 837,-2585 906,-2585 906,-2564 837,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"842\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 8&#45;&gt;28 -->\n",
+       "<g id=\"edge33\" class=\"edge\">\n",
+       "<title>8&#45;&gt;28</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M824,-2641.63C824,-2633.82 824,-2624.73 824,-2616.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"827.5,-2616.16 824,-2606.16 820.5,-2616.16 827.5,-2616.16\"/>\n",
+       "</g>\n",
+       "<!-- 11 -->\n",
+       "<g id=\"node12\" class=\"node\">\n",
+       "<title>11</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"706.5,-2684 557.5,-2684 557.5,-2642 706.5,-2642 706.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"558,-2642 558,-2684 608,-2684 608,-2642 558,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"563\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"564.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"608,-2663 608,-2684 656,-2684 656,-2663 608,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"618\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"656,-2663 656,-2684 707,-2684 707,-2663 656,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"664\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"608,-2642 608,-2663 656,-2663 656,-2642 608,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"613\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"656,-2642 656,-2663 707,-2663 707,-2642 656,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"661\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 10&#45;&gt;11 -->\n",
+       "<g id=\"edge13\" class=\"edge\">\n",
+       "<title>10&#45;&gt;11</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2719.63C632,-2711.82 632,-2702.73 632,-2694.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2694.16 632,-2684.16 628.5,-2694.16 635.5,-2694.16\"/>\n",
+       "</g>\n",
+       "<!-- 24 -->\n",
+       "<g id=\"node25\" class=\"node\">\n",
+       "<title>24</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"705,-2606 559,-2606 559,-2564 705,-2564 705,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"559,-2564 559,-2606 606,-2606 606,-2564 559,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"567\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"564\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-2585 606,-2606 654,-2606 654,-2585 606,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"616\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-2585 654,-2606 705,-2606 705,-2585 654,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-2564 606,-2585 654,-2585 654,-2564 606,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"611\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-2564 654,-2585 705,-2585 705,-2564 654,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 11&#45;&gt;24 -->\n",
+       "<g id=\"edge28\" class=\"edge\">\n",
+       "<title>11&#45;&gt;24</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2641.63C632,-2633.82 632,-2624.73 632,-2616.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2616.16 632,-2606.16 628.5,-2616.16 635.5,-2616.16\"/>\n",
+       "</g>\n",
+       "<!-- 13 -->\n",
+       "<g id=\"node14\" class=\"node\">\n",
+       "<title>13</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"167.5,-2684 24.5,-2684 24.5,-2642 167.5,-2642 167.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"25,-2642 25,-2684 75,-2684 75,-2642 25,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"30\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"31.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"75,-2663 75,-2684 123,-2684 123,-2663 75,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"85\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"123,-2663 123,-2684 168,-2684 168,-2663 123,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"128\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"75,-2642 75,-2663 123,-2663 123,-2642 75,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"80\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"123,-2642 123,-2663 168,-2663 168,-2642 123,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"128\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 12&#45;&gt;13 -->\n",
+       "<g id=\"edge15\" class=\"edge\">\n",
+       "<title>12&#45;&gt;13</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M76.97,-2719.63C79.67,-2711.73 82.82,-2702.53 85.77,-2693.9\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"89.18,-2694.75 89.1,-2684.16 82.55,-2692.49 89.18,-2694.75\"/>\n",
+       "</g>\n",
+       "<!-- 14 -->\n",
+       "<g id=\"node15\" class=\"node\">\n",
+       "<title>14</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"228,-2606 70,-2606 70,-2564 228,-2564 228,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"70,-2564 70,-2606 135,-2606 135,-2564 70,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"75\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"84\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"135,-2585 135,-2606 183,-2606 183,-2585 135,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"145\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"183,-2585 183,-2606 228,-2606 228,-2585 183,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"188\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"135,-2564 135,-2585 183,-2585 183,-2564 135,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"140\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"183,-2564 183,-2585 228,-2585 228,-2564 183,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"188\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 13&#45;&gt;14 -->\n",
+       "<g id=\"edge16\" class=\"edge\">\n",
+       "<title>13&#45;&gt;14</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M110.2,-2641.63C116.01,-2633.3 122.83,-2623.52 129.12,-2614.5\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"132.09,-2616.36 134.94,-2606.16 126.35,-2612.36 132.09,-2616.36\"/>\n",
+       "</g>\n",
+       "<!-- 15 -->\n",
+       "<g id=\"node16\" class=\"node\">\n",
+       "<title>15</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"338,-2060 198,-2060 198,-2018 338,-2018 338,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"198,-2018 198,-2060 245,-2060 245,-2018 198,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"210\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"203\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"245,-2039 245,-2060 293,-2060 293,-2039 245,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"255\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"293,-2039 293,-2060 338,-2060 338,-2039 293,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"298\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"245,-2018 245,-2039 293,-2039 293,-2018 245,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"250\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"293,-2018 293,-2039 338,-2039 338,-2018 293,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"298\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 14&#45;&gt;15 -->\n",
+       "<g id=\"edge17\" class=\"edge\">\n",
+       "<title>14&#45;&gt;15</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M174.69,-2563.82C205.77,-2536.92 254,-2486.14 254,-2430 254,-2430 254,-2430 254,-2194 254,-2151.01 259.48,-2101.55 263.58,-2070.53\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"267.08,-2070.75 264.96,-2060.37 260.15,-2069.81 267.08,-2070.75\"/>\n",
+       "</g>\n",
+       "<!-- 45 -->\n",
+       "<g id=\"node46\" class=\"node\">\n",
+       "<title>45</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"745.5,-1748 518.5,-1748 518.5,-1706 745.5,-1706 745.5,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"519,-1706 519,-1748 566,-1748 566,-1706 519,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"534.5\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">cat</text>\n",
+       "<text text-anchor=\"start\" x=\"524\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"566,-1727 566,-1748 614,-1748 614,-1727 566,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"576\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"614,-1727 614,-1748 746,-1748 746,-1727 614,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"619\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10), 2 x (1, 32), (1, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"566,-1706 566,-1727 614,-1727 614,-1706 566,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"571\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"614,-1706 614,-1727 746,-1727 746,-1706 614,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"662.5\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "</g>\n",
+       "<!-- 15&#45;&gt;45 -->\n",
+       "<g id=\"edge18\" class=\"edge\">\n",
+       "<title>15&#45;&gt;45</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M271.92,-2017.64C282.59,-1968.26 317.1,-1841.92 400,-1784 431.73,-1761.83 471.18,-1748.39 508.51,-1740.25\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"509.37,-1743.65 518.46,-1738.19 507.96,-1736.79 509.37,-1743.65\"/>\n",
+       "</g>\n",
+       "<!-- 17 -->\n",
+       "<g id=\"node18\" class=\"node\">\n",
+       "<title>17</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"522.5,-2528 361.5,-2528 361.5,-2486 522.5,-2486 522.5,-2528\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"362,-2486 362,-2528 412,-2528 412,-2486 362,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"367\" y=\"-2510\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"368.5\" y=\"-2499\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"412,-2507 412,-2528 460,-2528 460,-2507 412,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"422\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-2507 460,-2528 523,-2528 523,-2507 460,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"412,-2486 412,-2507 460,-2507 460,-2486 412,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"417\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-2486 460,-2507 523,-2507 523,-2486 460,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 16&#45;&gt;17 -->\n",
+       "<g id=\"edge20\" class=\"edge\">\n",
+       "<title>16&#45;&gt;17</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M441.27,-2563.63C441.37,-2555.82 441.49,-2546.73 441.6,-2538.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"445.1,-2538.2 441.73,-2528.16 438.1,-2538.11 445.1,-2538.2\"/>\n",
+       "</g>\n",
+       "<!-- 18 -->\n",
+       "<g id=\"node19\" class=\"node\">\n",
+       "<title>18</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"531,-2450 355,-2450 355,-2408 531,-2408 531,-2450\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"355,-2408 355,-2450 420,-2450 420,-2408 355,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"360\" y=\"-2432\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"369\" y=\"-2421\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"420,-2429 420,-2450 468,-2450 468,-2429 420,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"430\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"468,-2429 468,-2450 531,-2450 531,-2429 468,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"420,-2408 420,-2429 468,-2429 468,-2408 420,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"425\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"468,-2408 468,-2429 531,-2429 531,-2408 468,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"473\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 17&#45;&gt;18 -->\n",
+       "<g id=\"edge21\" class=\"edge\">\n",
+       "<title>17&#45;&gt;18</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M442.27,-2485.63C442.37,-2477.82 442.49,-2468.73 442.6,-2460.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"446.1,-2460.2 442.73,-2450.16 439.1,-2460.11 446.1,-2460.2\"/>\n",
+       "</g>\n",
+       "<!-- 19 -->\n",
+       "<g id=\"node20\" class=\"node\">\n",
+       "<title>19</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"523,-2372 365,-2372 365,-2330 523,-2330 523,-2372\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"365,-2330 365,-2372 412,-2372 412,-2330 365,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"377\" y=\"-2354\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"370\" y=\"-2343\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"412,-2351 412,-2372 460,-2372 460,-2351 412,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"422\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-2351 460,-2372 523,-2372 523,-2351 460,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"412,-2330 412,-2351 460,-2351 460,-2330 412,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"417\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-2330 460,-2351 523,-2351 523,-2330 460,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 18&#45;&gt;19 -->\n",
+       "<g id=\"edge22\" class=\"edge\">\n",
+       "<title>18&#45;&gt;19</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M443.27,-2407.63C443.37,-2399.82 443.49,-2390.73 443.6,-2382.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"447.1,-2382.2 443.73,-2372.16 440.1,-2382.11 447.1,-2382.2\"/>\n",
+       "</g>\n",
+       "<!-- 20 -->\n",
+       "<g id=\"node21\" class=\"node\">\n",
+       "<title>20</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"525,-2294 367,-2294 367,-2252 525,-2252 525,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"367,-2252 367,-2294 414,-2294 414,-2252 367,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"375\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"372\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"414,-2273 414,-2294 462,-2294 462,-2273 414,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"424\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"462,-2273 462,-2294 525,-2294 525,-2273 462,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"467\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"414,-2252 414,-2273 462,-2273 462,-2252 414,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"419\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"462,-2252 462,-2273 525,-2273 525,-2252 462,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"467\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 19&#45;&gt;20 -->\n",
+       "<g id=\"edge23\" class=\"edge\">\n",
+       "<title>19&#45;&gt;20</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M444.54,-2329.63C444.74,-2321.82 444.98,-2312.73 445.21,-2304.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"448.71,-2304.25 445.47,-2294.16 441.71,-2304.06 448.71,-2304.25\"/>\n",
+       "</g>\n",
+       "<!-- 21 -->\n",
+       "<g id=\"node22\" class=\"node\">\n",
+       "<title>21</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"527.5,-2216 366.5,-2216 366.5,-2174 527.5,-2174 527.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"367,-2174 367,-2216 417,-2216 417,-2174 367,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"372\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"373.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"417,-2195 417,-2216 465,-2216 465,-2195 417,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"427\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-2195 465,-2216 528,-2216 528,-2195 465,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"417,-2174 417,-2195 465,-2195 465,-2174 417,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"422\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"465,-2174 465,-2195 528,-2195 528,-2174 465,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 20&#45;&gt;21 -->\n",
+       "<g id=\"edge24\" class=\"edge\">\n",
+       "<title>20&#45;&gt;21</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M446.27,-2251.63C446.37,-2243.82 446.49,-2234.73 446.6,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"450.1,-2226.2 446.73,-2216.16 443.1,-2226.11 450.1,-2226.2\"/>\n",
+       "</g>\n",
+       "<!-- 22 -->\n",
+       "<g id=\"node23\" class=\"node\">\n",
+       "<title>22</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"535,-2138 359,-2138 359,-2096 535,-2096 535,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"359,-2096 359,-2138 424,-2138 424,-2096 359,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"364\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"373\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"424,-2117 424,-2138 472,-2138 472,-2117 424,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"434\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"472,-2117 472,-2138 535,-2138 535,-2117 472,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"477\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"424,-2096 424,-2117 472,-2117 472,-2096 424,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"429\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"472,-2096 472,-2117 535,-2117 535,-2096 472,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"477\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 21&#45;&gt;22 -->\n",
+       "<g id=\"edge25\" class=\"edge\">\n",
+       "<title>21&#45;&gt;22</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M447,-2173.63C447,-2165.82 447,-2156.73 447,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"450.5,-2148.16 447,-2138.16 443.5,-2148.16 450.5,-2148.16\"/>\n",
+       "</g>\n",
+       "<!-- 23 -->\n",
+       "<g id=\"node24\" class=\"node\">\n",
+       "<title>23</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"527,-2060 369,-2060 369,-2018 527,-2018 527,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"369,-2018 369,-2060 416,-2060 416,-2018 369,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"381\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"374\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"416,-2039 416,-2060 464,-2060 464,-2039 416,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"426\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"464,-2039 464,-2060 527,-2060 527,-2039 464,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"469\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"416,-2018 416,-2039 464,-2039 464,-2018 416,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"421\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"464,-2018 464,-2039 527,-2039 527,-2018 464,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"469\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 22&#45;&gt;23 -->\n",
+       "<g id=\"edge26\" class=\"edge\">\n",
+       "<title>22&#45;&gt;23</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M447.27,-2095.63C447.37,-2087.82 447.49,-2078.73 447.6,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"451.1,-2070.2 447.73,-2060.16 444.1,-2070.11 451.1,-2070.2\"/>\n",
+       "</g>\n",
+       "<!-- 36 -->\n",
+       "<g id=\"node37\" class=\"node\">\n",
+       "<title>36</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"529.5,-1982 366.5,-1982 366.5,-1940 529.5,-1940 529.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"367,-1940 367,-1982 419,-1982 419,-1940 367,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"372\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"374.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"419,-1961 419,-1982 467,-1982 467,-1961 419,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"429\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"467,-1961 467,-1982 530,-1982 530,-1961 467,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"472\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"419,-1940 419,-1961 467,-1961 467,-1940 419,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"424\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"467,-1940 467,-1961 530,-1961 530,-1940 467,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"472\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "</g>\n",
+       "<!-- 23&#45;&gt;36 -->\n",
+       "<g id=\"edge27\" class=\"edge\">\n",
+       "<title>23&#45;&gt;36</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M448,-2017.63C448,-2009.82 448,-2000.73 448,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"451.5,-1992.16 448,-1982.16 444.5,-1992.16 451.5,-1992.16\"/>\n",
+       "</g>\n",
+       "<!-- 25 -->\n",
+       "<g id=\"node26\" class=\"node\">\n",
+       "<title>25</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"706.5,-2528 557.5,-2528 557.5,-2486 706.5,-2486 706.5,-2528\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"558,-2486 558,-2528 608,-2528 608,-2486 558,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"563\" y=\"-2510\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"564.5\" y=\"-2499\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"608,-2507 608,-2528 656,-2528 656,-2507 608,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"618\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"656,-2507 656,-2528 707,-2528 707,-2507 656,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"661\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"608,-2486 608,-2507 656,-2507 656,-2486 608,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"613\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"656,-2486 656,-2507 707,-2507 707,-2486 656,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"661\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 24&#45;&gt;25 -->\n",
+       "<g id=\"edge29\" class=\"edge\">\n",
+       "<title>24&#45;&gt;25</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2563.63C632,-2555.82 632,-2546.73 632,-2538.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2538.16 632,-2528.16 628.5,-2538.16 635.5,-2538.16\"/>\n",
+       "</g>\n",
+       "<!-- 26 -->\n",
+       "<g id=\"node27\" class=\"node\">\n",
+       "<title>26</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"714,-2450 550,-2450 550,-2408 714,-2408 714,-2450\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"550,-2408 550,-2450 615,-2450 615,-2408 550,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"555\" y=\"-2432\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"564\" y=\"-2421\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2429 615,-2450 663,-2450 663,-2429 615,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"625\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2429 663,-2450 714,-2450 714,-2429 663,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"615,-2408 615,-2429 663,-2429 663,-2408 615,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"620\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"663,-2408 663,-2429 714,-2429 714,-2408 663,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"668\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 25&#45;&gt;26 -->\n",
+       "<g id=\"edge30\" class=\"edge\">\n",
+       "<title>25&#45;&gt;26</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2485.63C632,-2477.82 632,-2468.73 632,-2460.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2460.16 632,-2450.16 628.5,-2460.16 635.5,-2460.16\"/>\n",
+       "</g>\n",
+       "<!-- 27 -->\n",
+       "<g id=\"node28\" class=\"node\">\n",
+       "<title>27</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"705,-2372 559,-2372 559,-2330 705,-2330 705,-2372\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"559,-2330 559,-2372 606,-2372 606,-2330 559,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"571\" y=\"-2354\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"564\" y=\"-2343\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-2351 606,-2372 654,-2372 654,-2351 606,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"616\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-2351 654,-2372 705,-2372 705,-2351 654,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"606,-2330 606,-2351 654,-2351 654,-2330 606,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"611\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"654,-2330 654,-2351 705,-2351 705,-2330 654,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"659\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 26&#45;&gt;27 -->\n",
+       "<g id=\"edge31\" class=\"edge\">\n",
+       "<title>26&#45;&gt;27</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2407.63C632,-2399.82 632,-2390.73 632,-2382.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2382.16 632,-2372.16 628.5,-2382.16 635.5,-2382.16\"/>\n",
+       "</g>\n",
+       "<!-- 39 -->\n",
+       "<g id=\"node40\" class=\"node\">\n",
+       "<title>39</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"707.5,-2294 556.5,-2294 556.5,-2252 707.5,-2252 707.5,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"557,-2252 557,-2294 609,-2294 609,-2252 557,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"562\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"564.5\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"609,-2273 609,-2294 657,-2294 657,-2273 609,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"619\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"657,-2273 657,-2294 708,-2294 708,-2273 657,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"662\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"609,-2252 609,-2273 657,-2273 657,-2252 609,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"614\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"657,-2252 657,-2273 708,-2273 708,-2252 657,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"662\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "</g>\n",
+       "<!-- 27&#45;&gt;39 -->\n",
+       "<g id=\"edge32\" class=\"edge\">\n",
+       "<title>27&#45;&gt;39</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2329.63C632,-2321.82 632,-2312.73 632,-2304.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2304.16 632,-2294.16 628.5,-2304.16 635.5,-2304.16\"/>\n",
+       "</g>\n",
+       "<!-- 29 -->\n",
+       "<g id=\"node30\" class=\"node\">\n",
+       "<title>29</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"907.5,-2528 740.5,-2528 740.5,-2486 907.5,-2486 907.5,-2528\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"741,-2486 741,-2528 791,-2528 791,-2486 741,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"746\" y=\"-2510\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"747.5\" y=\"-2499\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-2507 791,-2528 839,-2528 839,-2507 791,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"801\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-2507 839,-2528 908,-2528 908,-2507 839,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"791,-2486 791,-2507 839,-2507 839,-2486 791,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"796\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"839,-2486 839,-2507 908,-2507 908,-2486 839,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"844\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 28&#45;&gt;29 -->\n",
+       "<g id=\"edge34\" class=\"edge\">\n",
+       "<title>28&#45;&gt;29</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M824,-2563.63C824,-2555.82 824,-2546.73 824,-2538.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"827.5,-2538.16 824,-2528.16 820.5,-2538.16 827.5,-2538.16\"/>\n",
+       "</g>\n",
+       "<!-- 30 -->\n",
+       "<g id=\"node31\" class=\"node\">\n",
+       "<title>30</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"915,-2450 733,-2450 733,-2408 915,-2408 915,-2450\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"733,-2408 733,-2450 798,-2450 798,-2408 733,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"738\" y=\"-2432\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"747\" y=\"-2421\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"798,-2429 798,-2450 846,-2450 846,-2429 798,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"808\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"846,-2429 846,-2450 915,-2450 915,-2429 846,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"851\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"798,-2408 798,-2429 846,-2429 846,-2408 798,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"803\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"846,-2408 846,-2429 915,-2429 915,-2408 846,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"851\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 29&#45;&gt;30 -->\n",
+       "<g id=\"edge35\" class=\"edge\">\n",
+       "<title>29&#45;&gt;30</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M824,-2485.63C824,-2477.82 824,-2468.73 824,-2460.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"827.5,-2460.16 824,-2450.16 820.5,-2460.16 827.5,-2460.16\"/>\n",
+       "</g>\n",
+       "<!-- 31 -->\n",
+       "<g id=\"node32\" class=\"node\">\n",
+       "<title>31</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"905,-2372 741,-2372 741,-2330 905,-2330 905,-2372\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"741,-2330 741,-2372 788,-2372 788,-2330 741,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"753\" y=\"-2354\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"746\" y=\"-2343\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"788,-2351 788,-2372 836,-2372 836,-2351 788,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"798\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"836,-2351 836,-2372 905,-2372 905,-2351 836,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"841\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"788,-2330 788,-2351 836,-2351 836,-2330 788,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"793\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"836,-2330 836,-2351 905,-2351 905,-2330 836,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"841\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 30&#45;&gt;31 -->\n",
+       "<g id=\"edge36\" class=\"edge\">\n",
+       "<title>30&#45;&gt;31</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M823.73,-2407.63C823.63,-2399.82 823.51,-2390.73 823.4,-2382.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"826.9,-2382.11 823.27,-2372.16 819.9,-2382.2 826.9,-2382.11\"/>\n",
+       "</g>\n",
+       "<!-- 32 -->\n",
+       "<g id=\"node33\" class=\"node\">\n",
+       "<title>32</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"903,-2294 739,-2294 739,-2252 903,-2252 903,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"739,-2252 739,-2294 786,-2294 786,-2252 739,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"747\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"744\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"786,-2273 786,-2294 834,-2294 834,-2273 786,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"796\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"834,-2273 834,-2294 903,-2294 903,-2273 834,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"839\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"786,-2252 786,-2273 834,-2273 834,-2252 786,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"791\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"834,-2252 834,-2273 903,-2273 903,-2252 834,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"839\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 31&#45;&gt;32 -->\n",
+       "<g id=\"edge37\" class=\"edge\">\n",
+       "<title>31&#45;&gt;32</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M822.46,-2329.63C822.26,-2321.82 822.02,-2312.73 821.79,-2304.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"825.29,-2304.06 821.53,-2294.16 818.29,-2304.25 825.29,-2304.06\"/>\n",
+       "</g>\n",
+       "<!-- 33 -->\n",
+       "<g id=\"node34\" class=\"node\">\n",
+       "<title>33</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"903.5,-2216 736.5,-2216 736.5,-2174 903.5,-2174 903.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"737,-2174 737,-2216 787,-2216 787,-2174 737,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"742\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"743.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-2195 787,-2216 835,-2216 835,-2195 787,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"797\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-2195 835,-2216 904,-2216 904,-2195 835,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-2174 787,-2195 835,-2195 835,-2174 787,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"792\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-2174 835,-2195 904,-2195 904,-2174 835,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 32&#45;&gt;33 -->\n",
+       "<g id=\"edge38\" class=\"edge\">\n",
+       "<title>32&#45;&gt;33</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M820.73,-2251.63C820.63,-2243.82 820.51,-2234.73 820.4,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"823.9,-2226.11 820.27,-2216.16 816.9,-2226.2 823.9,-2226.11\"/>\n",
+       "</g>\n",
+       "<!-- 34 -->\n",
+       "<g id=\"node35\" class=\"node\">\n",
+       "<title>34</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"911,-2138 729,-2138 729,-2096 911,-2096 911,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"729,-2096 729,-2138 794,-2138 794,-2096 729,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"734\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"743\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"794,-2117 794,-2138 842,-2138 842,-2117 794,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"804\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"842,-2117 842,-2138 911,-2138 911,-2117 842,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"847\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"794,-2096 794,-2117 842,-2117 842,-2096 794,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"799\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"842,-2096 842,-2117 911,-2117 911,-2096 842,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"847\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 33&#45;&gt;34 -->\n",
+       "<g id=\"edge39\" class=\"edge\">\n",
+       "<title>33&#45;&gt;34</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M820,-2173.63C820,-2165.82 820,-2156.73 820,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"823.5,-2148.16 820,-2138.16 816.5,-2148.16 823.5,-2148.16\"/>\n",
+       "</g>\n",
+       "<!-- 35 -->\n",
+       "<g id=\"node36\" class=\"node\">\n",
+       "<title>35</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"901,-2060 737,-2060 737,-2018 901,-2018 901,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"737,-2018 737,-2060 784,-2060 784,-2018 737,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"749\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"742\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"784,-2039 784,-2060 832,-2060 832,-2039 784,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"794\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"832,-2039 832,-2060 901,-2060 901,-2039 832,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"837\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"784,-2018 784,-2039 832,-2039 832,-2018 784,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"789\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"832,-2018 832,-2039 901,-2039 901,-2018 832,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"837\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 34&#45;&gt;35 -->\n",
+       "<g id=\"edge40\" class=\"edge\">\n",
+       "<title>34&#45;&gt;35</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M819.73,-2095.63C819.63,-2087.82 819.51,-2078.73 819.4,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"822.9,-2070.11 819.27,-2060.16 815.9,-2070.2 822.9,-2070.11\"/>\n",
+       "</g>\n",
+       "<!-- 42 -->\n",
+       "<g id=\"node43\" class=\"node\">\n",
+       "<title>42</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"903.5,-1982 734.5,-1982 734.5,-1940 903.5,-1940 903.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"735,-1940 735,-1982 787,-1982 787,-1940 735,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"740\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"742.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-1961 787,-1982 835,-1982 835,-1961 787,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"797\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-1961 835,-1982 904,-1982 904,-1961 835,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"787,-1940 787,-1961 835,-1961 835,-1940 787,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"792\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"835,-1940 835,-1961 904,-1961 904,-1940 835,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"840\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "</g>\n",
+       "<!-- 35&#45;&gt;42 -->\n",
+       "<g id=\"edge41\" class=\"edge\">\n",
+       "<title>35&#45;&gt;42</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M819,-2017.63C819,-2009.82 819,-2000.73 819,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"822.5,-1992.16 819,-1982.16 815.5,-1992.16 822.5,-1992.16\"/>\n",
+       "</g>\n",
+       "<!-- 37 -->\n",
+       "<g id=\"node38\" class=\"node\">\n",
+       "<title>37</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"539,-1904 357,-1904 357,-1862 539,-1862 539,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"357,-1862 357,-1904 428,-1904 428,-1862 357,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"362\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"374\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"428,-1883 428,-1904 476,-1904 476,-1883 428,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"438\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"476,-1883 476,-1904 539,-1904 539,-1883 476,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"481\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"428,-1862 428,-1883 476,-1883 476,-1862 428,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"433\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"476,-1862 476,-1883 539,-1883 539,-1862 476,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"484\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 36&#45;&gt;37 -->\n",
+       "<g id=\"edge42\" class=\"edge\">\n",
+       "<title>36&#45;&gt;37</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M448,-1939.63C448,-1931.82 448,-1922.73 448,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"451.5,-1914.16 448,-1904.16 444.5,-1914.16 451.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 38 -->\n",
+       "<g id=\"node39\" class=\"node\">\n",
+       "<title>38</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"565,-1826 409,-1826 409,-1784 565,-1784 565,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"409,-1784 409,-1826 460,-1826 460,-1784 409,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"414\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"416\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-1805 460,-1826 508,-1826 508,-1805 460,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"470\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"508,-1805 508,-1826 565,-1826 565,-1805 508,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"513\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"460,-1784 460,-1805 508,-1805 508,-1784 460,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"465\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"508,-1784 508,-1805 565,-1805 565,-1784 508,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"519\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 37&#45;&gt;38 -->\n",
+       "<g id=\"edge43\" class=\"edge\">\n",
+       "<title>37&#45;&gt;38</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M458.45,-1861.63C462.64,-1853.48 467.54,-1843.92 472.09,-1835.06\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"475.2,-1836.65 476.66,-1826.16 468.98,-1833.46 475.2,-1836.65\"/>\n",
+       "</g>\n",
+       "<!-- 38&#45;&gt;45 -->\n",
+       "<g id=\"edge44\" class=\"edge\">\n",
+       "<title>38&#45;&gt;45</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M525.48,-1783.83C543.66,-1774.3 565.53,-1762.84 584.71,-1752.79\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"586.36,-1755.87 593.6,-1748.13 583.11,-1749.67 586.36,-1755.87\"/>\n",
+       "</g>\n",
+       "<!-- 40 -->\n",
+       "<g id=\"node41\" class=\"node\">\n",
+       "<title>40</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"717,-2216 547,-2216 547,-2174 717,-2174 717,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"547,-2174 547,-2216 618,-2216 618,-2174 547,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"552\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"564\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"618,-2195 618,-2216 666,-2216 666,-2195 618,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"628\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"666,-2195 666,-2216 717,-2216 717,-2195 666,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"671\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"618,-2174 618,-2195 666,-2195 666,-2174 618,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"623\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"666,-2174 666,-2195 717,-2195 717,-2174 666,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"671\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "</g>\n",
+       "<!-- 39&#45;&gt;40 -->\n",
+       "<g id=\"edge45\" class=\"edge\">\n",
+       "<title>39&#45;&gt;40</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2251.63C632,-2243.82 632,-2234.73 632,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-2226.16 632,-2216.16 628.5,-2226.16 635.5,-2226.16\"/>\n",
+       "</g>\n",
+       "<!-- 41 -->\n",
+       "<g id=\"node42\" class=\"node\">\n",
+       "<title>41</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"707,-1904 557,-1904 557,-1862 707,-1862 707,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"557,-1862 557,-1904 608,-1904 608,-1862 557,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"562\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"564\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"608,-1883 608,-1904 656,-1904 656,-1883 608,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"618\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"656,-1883 656,-1904 707,-1904 707,-1883 656,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"661\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"608,-1862 608,-1883 656,-1883 656,-1862 608,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"613\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"656,-1862 656,-1883 707,-1883 707,-1862 656,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"667\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3) </text>\n",
+       "</g>\n",
+       "<!-- 40&#45;&gt;41 -->\n",
+       "<g id=\"edge46\" class=\"edge\">\n",
+       "<title>40&#45;&gt;41</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-2173.66C632,-2120.43 632,-1976.84 632,-1914.06\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-1914.01 632,-1904.01 628.5,-1914.01 635.5,-1914.01\"/>\n",
+       "</g>\n",
+       "<!-- 41&#45;&gt;45 -->\n",
+       "<g id=\"edge47\" class=\"edge\">\n",
+       "<title>41&#45;&gt;45</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632,-1861.87C632,-1835.47 632,-1788.91 632,-1758.26\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"635.5,-1758.14 632,-1748.14 628.5,-1758.14 635.5,-1758.14\"/>\n",
+       "</g>\n",
+       "<!-- 43 -->\n",
+       "<g id=\"node44\" class=\"node\">\n",
+       "<title>43</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"913,-1904 725,-1904 725,-1862 913,-1862 913,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"725,-1862 725,-1904 796,-1904 796,-1862 725,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"730\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"742\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"796,-1883 796,-1904 844,-1904 844,-1883 796,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"806\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"844,-1883 844,-1904 913,-1904 913,-1883 844,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"849\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"796,-1862 796,-1883 844,-1883 844,-1862 796,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"801\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"844,-1862 844,-1883 913,-1883 913,-1862 844,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"855\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 42&#45;&gt;43 -->\n",
+       "<g id=\"edge48\" class=\"edge\">\n",
+       "<title>42&#45;&gt;43</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M819,-1939.63C819,-1931.82 819,-1922.73 819,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"822.5,-1914.16 819,-1904.16 815.5,-1914.16 822.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 44 -->\n",
+       "<g id=\"node45\" class=\"node\">\n",
+       "<title>44</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"857,-1826 701,-1826 701,-1784 857,-1784 857,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"701,-1784 701,-1826 752,-1826 752,-1784 701,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"706\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"708\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"752,-1805 752,-1826 800,-1826 800,-1805 752,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"762\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"800,-1805 800,-1826 857,-1826 857,-1805 800,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"805\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"752,-1784 752,-1805 800,-1805 800,-1784 752,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"757\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"800,-1784 800,-1805 857,-1805 857,-1784 800,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"811\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 43&#45;&gt;44 -->\n",
+       "<g id=\"edge49\" class=\"edge\">\n",
+       "<title>43&#45;&gt;44</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M808.28,-1861.63C803.99,-1853.48 798.96,-1843.92 794.29,-1835.06\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"797.36,-1833.38 789.61,-1826.16 791.17,-1836.64 797.36,-1833.38\"/>\n",
+       "</g>\n",
+       "<!-- 44&#45;&gt;45 -->\n",
+       "<g id=\"edge50\" class=\"edge\">\n",
+       "<title>44&#45;&gt;45</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M739.99,-1783.83C721.56,-1774.3 699.39,-1762.84 679.94,-1752.79\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"681.42,-1749.61 670.93,-1748.13 678.21,-1755.83 681.42,-1749.61\"/>\n",
+       "</g>\n",
+       "<!-- 46 -->\n",
+       "<g id=\"node47\" class=\"node\">\n",
+       "<title>46</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"706,-1670 560,-1670 560,-1628 706,-1628 706,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"560,-1628 560,-1670 607,-1670 607,-1628 560,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"568\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"565\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"607,-1649 607,-1670 655,-1670 655,-1649 607,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"617\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"655,-1649 655,-1670 706,-1670 706,-1649 655,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"663\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"607,-1628 607,-1649 655,-1649 655,-1628 607,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"612\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"655,-1628 655,-1649 706,-1649 706,-1628 655,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"660\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 45&#45;&gt;46 -->\n",
+       "<g id=\"edge51\" class=\"edge\">\n",
+       "<title>45&#45;&gt;46</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M632.27,-1705.63C632.37,-1697.82 632.49,-1688.73 632.6,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"636.1,-1680.2 632.73,-1670.16 629.1,-1680.11 636.1,-1680.2\"/>\n",
+       "</g>\n",
+       "<!-- 47 -->\n",
+       "<g id=\"node48\" class=\"node\">\n",
+       "<title>47</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"906,-1592 742,-1592 742,-1550 906,-1550 906,-1592\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"742,-1550 742,-1592 807,-1592 807,-1550 742,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"747\" y=\"-1574\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"756\" y=\"-1563\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"807,-1571 807,-1592 855,-1592 855,-1571 807,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"817\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"855,-1571 855,-1592 906,-1592 906,-1571 855,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"860\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"807,-1550 807,-1571 855,-1571 855,-1550 807,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"812\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"855,-1550 855,-1571 906,-1571 906,-1550 855,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"860\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 46&#45;&gt;47 -->\n",
+       "<g id=\"edge52\" class=\"edge\">\n",
+       "<title>46&#45;&gt;47</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M683.43,-1627.93C708.24,-1618.06 738.31,-1606.1 764.26,-1595.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"765.62,-1599 773.62,-1592.05 763.03,-1592.49 765.62,-1599\"/>\n",
+       "</g>\n",
+       "<!-- 48 -->\n",
+       "<g id=\"node49\" class=\"node\">\n",
+       "<title>48</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"993,-1514 847,-1514 847,-1472 993,-1472 993,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"847,-1472 847,-1514 894,-1514 894,-1472 847,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"859\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"852\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"894,-1493 894,-1514 942,-1514 942,-1493 894,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"904\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"942,-1493 942,-1514 993,-1514 993,-1493 942,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"947\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"894,-1472 894,-1493 942,-1493 942,-1472 894,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"899\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"942,-1472 942,-1493 993,-1493 993,-1472 942,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"947\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 47&#45;&gt;48 -->\n",
+       "<g id=\"edge53\" class=\"edge\">\n",
+       "<title>47&#45;&gt;48</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M849.47,-1549.83C860.76,-1540.9 874.19,-1530.27 886.29,-1520.69\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"888.74,-1523.21 894.41,-1514.26 884.39,-1517.72 888.74,-1523.21\"/>\n",
+       "</g>\n",
+       "<!-- 49 -->\n",
+       "<g id=\"node50\" class=\"node\">\n",
+       "<title>49</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1042.5,-1436 893.5,-1436 893.5,-1394 1042.5,-1394 1042.5,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"894,-1394 894,-1436 944,-1436 944,-1394 894,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"899\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"900.5\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-1415 944,-1436 992,-1436 992,-1415 944,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"954\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-1415 992,-1436 1043,-1436 1043,-1415 992,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"997\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-1394 944,-1415 992,-1415 992,-1394 944,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"949\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-1394 992,-1415 1043,-1415 1043,-1394 992,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"997\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 48&#45;&gt;49 -->\n",
+       "<g id=\"edge54\" class=\"edge\">\n",
+       "<title>48&#45;&gt;49</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M932.86,-1471.63C938.07,-1463.39 944.18,-1453.72 949.82,-1444.78\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"952.89,-1446.48 955.27,-1436.16 946.97,-1442.75 952.89,-1446.48\"/>\n",
+       "</g>\n",
+       "<!-- 50 -->\n",
+       "<g id=\"node51\" class=\"node\">\n",
+       "<title>50</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1065,-1358 919,-1358 919,-1316 1065,-1316 1065,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"919,-1316 919,-1358 966,-1358 966,-1316 919,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"927\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"924\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"966,-1337 966,-1358 1014,-1358 1014,-1337 966,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"976\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1014,-1337 1014,-1358 1065,-1358 1065,-1337 1014,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1019\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"966,-1316 966,-1337 1014,-1337 1014,-1316 966,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"971\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1014,-1316 1014,-1337 1065,-1337 1065,-1316 1014,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1019\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 49&#45;&gt;50 -->\n",
+       "<g id=\"edge55\" class=\"edge\">\n",
+       "<title>49&#45;&gt;50</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M974.43,-1393.63C976.93,-1385.73 979.83,-1376.53 982.56,-1367.9\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"985.96,-1368.75 985.63,-1358.16 979.28,-1366.64 985.96,-1368.75\"/>\n",
+       "</g>\n",
+       "<!-- 51 -->\n",
+       "<g id=\"node52\" class=\"node\">\n",
+       "<title>51</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1086,-1280 922,-1280 922,-1238 1086,-1238 1086,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"922,-1238 922,-1280 987,-1280 987,-1238 922,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"927\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"936\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"987,-1259 987,-1280 1035,-1280 1035,-1259 987,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"997\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1035,-1259 1035,-1280 1086,-1280 1086,-1259 1035,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1040\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"987,-1238 987,-1259 1035,-1259 1035,-1238 987,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"992\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1035,-1238 1035,-1259 1086,-1259 1086,-1238 1035,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1040\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 50&#45;&gt;51 -->\n",
+       "<g id=\"edge56\" class=\"edge\">\n",
+       "<title>50&#45;&gt;51</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M995.22,-1315.63C996.45,-1307.82 997.89,-1298.73 999.23,-1290.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1002.71,-1290.58 1000.82,-1280.16 995.8,-1289.49 1002.71,-1290.58\"/>\n",
+       "</g>\n",
+       "<!-- 52 -->\n",
+       "<g id=\"node53\" class=\"node\">\n",
+       "<title>52</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1083,-1202 937,-1202 937,-1160 1083,-1160 1083,-1202\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"937,-1160 937,-1202 984,-1202 984,-1160 937,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"949\" y=\"-1184\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"942\" y=\"-1173\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"984,-1181 984,-1202 1032,-1202 1032,-1181 984,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"994\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1032,-1181 1032,-1202 1083,-1202 1083,-1181 1032,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"1037\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"984,-1160 984,-1181 1032,-1181 1032,-1160 984,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"989\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1032,-1160 1032,-1181 1083,-1181 1083,-1160 1032,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1037\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 51&#45;&gt;52 -->\n",
+       "<g id=\"edge57\" class=\"edge\">\n",
+       "<title>51&#45;&gt;52</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1005.61,-1237.63C1006.22,-1229.82 1006.94,-1220.73 1007.62,-1212.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1011.11,-1212.4 1008.41,-1202.16 1004.13,-1211.85 1011.11,-1212.4\"/>\n",
+       "</g>\n",
+       "<!-- 53 -->\n",
+       "<g id=\"node54\" class=\"node\">\n",
+       "<title>53</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1087.5,-1124 938.5,-1124 938.5,-1082 1087.5,-1082 1087.5,-1124\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"939,-1082 939,-1124 989,-1124 989,-1082 939,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"944\" y=\"-1106\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"945.5\" y=\"-1095\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"989,-1103 989,-1124 1037,-1124 1037,-1103 989,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"999\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1037,-1103 1037,-1124 1088,-1124 1088,-1103 1037,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"1042\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"989,-1082 989,-1103 1037,-1103 1037,-1082 989,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"994\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1037,-1082 1037,-1103 1088,-1103 1088,-1082 1037,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1042\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 52&#45;&gt;53 -->\n",
+       "<g id=\"edge58\" class=\"edge\">\n",
+       "<title>52&#45;&gt;53</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1010.8,-1159.63C1011.11,-1151.82 1011.47,-1142.73 1011.81,-1134.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1015.31,-1134.29 1012.2,-1124.16 1008.31,-1134.01 1015.31,-1134.29\"/>\n",
+       "</g>\n",
+       "<!-- 54 -->\n",
+       "<g id=\"node55\" class=\"node\">\n",
+       "<title>54</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1087,-1046 941,-1046 941,-1004 1087,-1004 1087,-1046\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"941,-1004 941,-1046 988,-1046 988,-1004 941,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"949\" y=\"-1028\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"946\" y=\"-1017\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"988,-1025 988,-1046 1036,-1046 1036,-1025 988,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"998\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1036,-1025 1036,-1046 1087,-1046 1087,-1025 1036,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"1041\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"988,-1004 988,-1025 1036,-1025 1036,-1004 988,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"993\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1036,-1004 1036,-1025 1087,-1025 1087,-1004 1036,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1044\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 53&#45;&gt;54 -->\n",
+       "<g id=\"edge59\" class=\"edge\">\n",
+       "<title>53&#45;&gt;54</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1013.27,-1081.63C1013.37,-1073.82 1013.49,-1064.73 1013.6,-1056.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1017.1,-1056.2 1013.73,-1046.16 1010.1,-1056.11 1017.1,-1056.2\"/>\n",
+       "</g>\n",
+       "<!-- 55 -->\n",
+       "<g id=\"node56\" class=\"node\">\n",
+       "<title>55</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1094,-968 936,-968 936,-926 1094,-926 1094,-968\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"936,-926 936,-968 1001,-968 1001,-926 936,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"941\" y=\"-950\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"950\" y=\"-939\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1001,-947 1001,-968 1049,-968 1049,-947 1001,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"1011\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1049,-947 1049,-968 1094,-968 1094,-947 1049,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"1054\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1001,-926 1001,-947 1049,-947 1049,-926 1001,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1006\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1049,-926 1049,-947 1094,-947 1094,-926 1049,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1054\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 54&#45;&gt;55 -->\n",
+       "<g id=\"edge60\" class=\"edge\">\n",
+       "<title>54&#45;&gt;55</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1014.27,-1003.63C1014.37,-995.82 1014.49,-986.73 1014.6,-978.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1018.1,-978.2 1014.73,-968.16 1011.1,-978.11 1018.1,-978.2\"/>\n",
+       "</g>\n",
+       "<!-- 56 -->\n",
+       "<g id=\"node57\" class=\"node\">\n",
+       "<title>56</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1085,-890 945,-890 945,-848 1085,-848 1085,-890\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"945,-848 945,-890 992,-890 992,-848 945,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"957\" y=\"-872\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"950\" y=\"-861\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-869 992,-890 1040,-890 1040,-869 992,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"1002\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1040,-869 1040,-890 1085,-890 1085,-869 1040,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"1045\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-848 992,-869 1040,-869 1040,-848 992,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"997\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1040,-848 1040,-869 1085,-869 1085,-848 1040,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1045\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 55&#45;&gt;56 -->\n",
+       "<g id=\"edge61\" class=\"edge\">\n",
+       "<title>55&#45;&gt;56</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1015,-925.63C1015,-917.82 1015,-908.73 1015,-900.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1018.5,-900.16 1015,-890.16 1011.5,-900.16 1018.5,-900.16\"/>\n",
+       "</g>\n",
+       "<!-- 57 -->\n",
+       "<g id=\"node58\" class=\"node\">\n",
+       "<title>57</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1086.5,-812 943.5,-812 943.5,-770 1086.5,-770 1086.5,-812\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"944,-770 944,-812 994,-812 994,-770 944,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"949\" y=\"-794\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"950.5\" y=\"-783\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"994,-791 994,-812 1042,-812 1042,-791 994,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"1004\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1042,-791 1042,-812 1087,-812 1087,-791 1042,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"1047\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"994,-770 994,-791 1042,-791 1042,-770 994,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"999\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1042,-770 1042,-791 1087,-791 1087,-770 1042,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1047\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 56&#45;&gt;57 -->\n",
+       "<g id=\"edge62\" class=\"edge\">\n",
+       "<title>56&#45;&gt;57</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1015,-847.63C1015,-839.82 1015,-830.73 1015,-822.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1018.5,-822.16 1015,-812.16 1011.5,-822.16 1018.5,-822.16\"/>\n",
+       "</g>\n",
+       "<!-- 58 -->\n",
+       "<g id=\"node59\" class=\"node\">\n",
+       "<title>58</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1085,-734 945,-734 945,-692 1085,-692 1085,-734\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"945,-692 945,-734 992,-734 992,-692 945,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"953\" y=\"-716\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"950\" y=\"-705\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-713 992,-734 1040,-734 1040,-713 992,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"1002\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1040,-713 1040,-734 1085,-734 1085,-713 1040,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"1045\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-692 992,-713 1040,-713 1040,-692 992,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"997\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1040,-692 1040,-713 1085,-713 1085,-692 1040,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1045\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 57&#45;&gt;58 -->\n",
+       "<g id=\"edge63\" class=\"edge\">\n",
+       "<title>57&#45;&gt;58</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1015,-769.63C1015,-761.82 1015,-752.73 1015,-744.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1018.5,-744.16 1015,-734.16 1011.5,-744.16 1018.5,-744.16\"/>\n",
+       "</g>\n",
+       "<!-- 59 -->\n",
+       "<g id=\"node60\" class=\"node\">\n",
+       "<title>59</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1094,-656 936,-656 936,-614 1094,-614 1094,-656\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"936,-614 936,-656 1001,-656 1001,-614 936,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"941\" y=\"-638\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"950\" y=\"-627\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1001,-635 1001,-656 1049,-656 1049,-635 1001,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1011\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1049,-635 1049,-656 1094,-656 1094,-635 1049,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1054\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1001,-614 1001,-635 1049,-635 1049,-614 1001,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1006\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1049,-614 1049,-635 1094,-635 1094,-614 1049,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1054\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 58&#45;&gt;59 -->\n",
+       "<g id=\"edge64\" class=\"edge\">\n",
+       "<title>58&#45;&gt;59</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1015,-691.63C1015,-683.82 1015,-674.73 1015,-666.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1018.5,-666.16 1015,-656.16 1011.5,-666.16 1018.5,-666.16\"/>\n",
+       "</g>\n",
+       "<!-- 119 -->\n",
+       "<g id=\"node120\" class=\"node\">\n",
+       "<title>119</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1085,-578 945,-578 945,-536 1085,-536 1085,-578\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"945,-536 945,-578 992,-578 992,-536 945,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"953\" y=\"-560\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"950\" y=\"-549\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-557 992,-578 1040,-578 1040,-557 992,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1002\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1040,-557 1040,-578 1085,-578 1085,-557 1040,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1045\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"992,-536 992,-557 1040,-557 1040,-536 992,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"997\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1040,-536 1040,-557 1085,-557 1085,-536 1040,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1045\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 59&#45;&gt;119 -->\n",
+       "<g id=\"edge65\" class=\"edge\">\n",
+       "<title>59&#45;&gt;119</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1015,-613.63C1015,-605.82 1015,-596.73 1015,-588.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1018.5,-588.16 1015,-578.16 1011.5,-588.16 1018.5,-588.16\"/>\n",
+       "</g>\n",
+       "<!-- 69 -->\n",
+       "<g id=\"node70\" class=\"node\">\n",
+       "<title>69</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1071,-2762 931,-2762 931,-2720 1071,-2720 1071,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"931,-2720 931,-2762 978,-2762 978,-2720 931,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"939\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"936\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"978,-2741 978,-2762 1026,-2762 1026,-2741 978,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"988\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1026,-2741 1026,-2762 1071,-2762 1071,-2741 1026,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1031\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"978,-2720 978,-2741 1026,-2741 1026,-2720 978,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"983\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1026,-2720 1026,-2741 1071,-2741 1071,-2720 1026,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1031\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 60&#45;&gt;69 -->\n",
+       "<g id=\"edge73\" class=\"edge\">\n",
+       "<title>60&#45;&gt;69</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M936.58,-2797.63C946.82,-2788.79 958.96,-2778.3 969.93,-2768.84\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"972.38,-2771.34 977.66,-2762.16 967.8,-2766.05 972.38,-2771.34\"/>\n",
+       "</g>\n",
+       "<!-- 62 -->\n",
+       "<g id=\"node63\" class=\"node\">\n",
+       "<title>62</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1457.5,-2762 1282.5,-2762 1282.5,-2720 1457.5,-2720 1457.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1283,-2720 1283,-2762 1353,-2762 1353,-2720 1283,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1288\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1299.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2741 1353,-2762 1401,-2762 1401,-2741 1353,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2741 1401,-2762 1458,-2762 1458,-2741 1401,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-2720 1353,-2741 1401,-2741 1401,-2720 1353,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-2720 1401,-2741 1458,-2741 1458,-2720 1401,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "</g>\n",
+       "<!-- 61&#45;&gt;62 -->\n",
+       "<g id=\"edge66\" class=\"edge\">\n",
+       "<title>61&#45;&gt;62</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2797.63C1370,-2789.82 1370,-2780.73 1370,-2772.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2772.16 1370,-2762.16 1366.5,-2772.16 1373.5,-2772.16\"/>\n",
+       "</g>\n",
+       "<!-- 64 -->\n",
+       "<g id=\"node65\" class=\"node\">\n",
+       "<title>64</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1843.5,-2762 1668.5,-2762 1668.5,-2720 1843.5,-2720 1843.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1669,-2720 1669,-2762 1739,-2762 1739,-2720 1669,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1674\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1685.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1739,-2741 1739,-2762 1787,-2762 1787,-2741 1739,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1749\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1787,-2741 1787,-2762 1844,-2762 1844,-2741 1787,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1792\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1739,-2720 1739,-2741 1787,-2741 1787,-2720 1739,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1744\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1787,-2720 1787,-2741 1844,-2741 1844,-2720 1787,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1792\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "</g>\n",
+       "<!-- 61&#45;&gt;64 -->\n",
+       "<g id=\"edge67\" class=\"edge\">\n",
+       "<title>61&#45;&gt;64</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1457.63,-2801.3C1512.5,-2790.78 1584.93,-2776.77 1658.35,-2762.12\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1659.12,-2765.54 1668.24,-2760.14 1657.75,-2758.67 1659.12,-2765.54\"/>\n",
+       "</g>\n",
+       "<!-- 66 -->\n",
+       "<g id=\"node67\" class=\"node\">\n",
+       "<title>66</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1264.5,-2762 1089.5,-2762 1089.5,-2720 1264.5,-2720 1264.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1090,-2720 1090,-2762 1160,-2762 1160,-2720 1090,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1095\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1106.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1160,-2741 1160,-2762 1208,-2762 1208,-2741 1160,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1170\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1208,-2741 1208,-2762 1265,-2762 1265,-2741 1208,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1213\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1160,-2720 1160,-2741 1208,-2741 1208,-2720 1160,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1165\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1208,-2720 1208,-2741 1265,-2741 1265,-2720 1208,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1222\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 0) </text>\n",
+       "</g>\n",
+       "<!-- 61&#45;&gt;66 -->\n",
+       "<g id=\"edge68\" class=\"edge\">\n",
+       "<title>61&#45;&gt;66</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1319.04,-2797.93C1293.97,-2788.06 1263.59,-2776.1 1237.37,-2765.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1238.5,-2762.46 1227.91,-2762.05 1235.93,-2768.97 1238.5,-2762.46\"/>\n",
+       "</g>\n",
+       "<!-- 67 -->\n",
+       "<g id=\"node68\" class=\"node\">\n",
+       "<title>67</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1650.5,-2762 1475.5,-2762 1475.5,-2720 1650.5,-2720 1650.5,-2762\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1476,-2720 1476,-2762 1546,-2762 1546,-2720 1476,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1481\" y=\"-2744\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1492.5\" y=\"-2733\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1546,-2741 1546,-2762 1594,-2762 1594,-2741 1546,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1556\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1594,-2741 1594,-2762 1651,-2762 1651,-2741 1594,-2741\"/>\n",
+       "<text text-anchor=\"start\" x=\"1599\" y=\"-2749\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6720) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1546,-2720 1546,-2741 1594,-2741 1594,-2720 1546,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1551\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1594,-2720 1594,-2741 1651,-2741 1651,-2720 1594,-2720\"/>\n",
+       "<text text-anchor=\"start\" x=\"1605\" y=\"-2728\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "</g>\n",
+       "<!-- 61&#45;&gt;67 -->\n",
+       "<g id=\"edge69\" class=\"edge\">\n",
+       "<title>61&#45;&gt;67</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1420.96,-2797.93C1446.03,-2788.06 1476.41,-2776.1 1502.63,-2765.77\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1504.07,-2768.97 1512.09,-2762.05 1501.5,-2762.46 1504.07,-2768.97\"/>\n",
+       "</g>\n",
+       "<!-- 63 -->\n",
+       "<g id=\"node64\" class=\"node\">\n",
+       "<title>63</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1450.5,-2684 1289.5,-2684 1289.5,-2642 1450.5,-2642 1450.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1290,-2642 1290,-2684 1340,-2684 1340,-2642 1290,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1295\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1296.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2663 1340,-2684 1388,-2684 1388,-2663 1340,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1350\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2663 1388,-2684 1451,-2684 1451,-2663 1388,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1399\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 208) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2642 1340,-2663 1388,-2663 1388,-2642 1340,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1345\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2642 1388,-2663 1451,-2663 1451,-2642 1388,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "</g>\n",
+       "<!-- 62&#45;&gt;63 -->\n",
+       "<g id=\"edge70\" class=\"edge\">\n",
+       "<title>62&#45;&gt;63</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2719.63C1370,-2711.82 1370,-2702.73 1370,-2694.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2694.16 1370,-2684.16 1366.5,-2694.16 1373.5,-2694.16\"/>\n",
+       "</g>\n",
+       "<!-- 73 -->\n",
+       "<g id=\"node74\" class=\"node\">\n",
+       "<title>73</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-2606 1291,-2606 1291,-2564 1449,-2564 1449,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-2564 1291,-2606 1338,-2606 1338,-2564 1291,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1299\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2585 1338,-2606 1386,-2606 1386,-2585 1338,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2585 1386,-2606 1449,-2606 1449,-2585 1386,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2564 1338,-2585 1386,-2585 1386,-2564 1338,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1343\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2564 1386,-2585 1449,-2585 1449,-2564 1386,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 63&#45;&gt;73 -->\n",
+       "<g id=\"edge78\" class=\"edge\">\n",
+       "<title>63&#45;&gt;73</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2641.63C1370,-2633.82 1370,-2624.73 1370,-2616.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2616.16 1370,-2606.16 1366.5,-2616.16 1373.5,-2616.16\"/>\n",
+       "</g>\n",
+       "<!-- 65 -->\n",
+       "<g id=\"node66\" class=\"node\">\n",
+       "<title>65</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1838.5,-2684 1671.5,-2684 1671.5,-2642 1838.5,-2642 1838.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1672,-2642 1672,-2684 1722,-2684 1722,-2642 1672,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1677\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1678.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1722,-2663 1722,-2684 1770,-2684 1770,-2663 1722,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1732\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1770,-2663 1770,-2684 1839,-2684 1839,-2663 1770,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1781\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 6500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1722,-2642 1722,-2663 1770,-2663 1770,-2642 1722,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1727\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1770,-2642 1770,-2663 1839,-2663 1839,-2642 1770,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1775\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "</g>\n",
+       "<!-- 64&#45;&gt;65 -->\n",
+       "<g id=\"edge71\" class=\"edge\">\n",
+       "<title>64&#45;&gt;65</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1755.73,-2719.63C1755.63,-2711.82 1755.51,-2702.73 1755.4,-2694.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1758.9,-2694.11 1755.27,-2684.16 1751.9,-2694.2 1758.9,-2694.11\"/>\n",
+       "</g>\n",
+       "<!-- 85 -->\n",
+       "<g id=\"node86\" class=\"node\">\n",
+       "<title>85</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1836,-2606 1672,-2606 1672,-2564 1836,-2564 1836,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1672,-2564 1672,-2606 1719,-2606 1719,-2564 1672,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1680\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1677\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1719,-2585 1719,-2606 1767,-2606 1767,-2585 1719,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1729\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1767,-2585 1767,-2606 1836,-2606 1836,-2585 1767,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1772\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 13) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1719,-2564 1719,-2585 1767,-2585 1767,-2564 1719,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1724\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1767,-2564 1767,-2585 1836,-2585 1836,-2564 1767,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1772\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 65&#45;&gt;85 -->\n",
+       "<g id=\"edge92\" class=\"edge\">\n",
+       "<title>65&#45;&gt;85</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1754.73,-2641.63C1754.63,-2633.82 1754.51,-2624.73 1754.4,-2616.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1757.9,-2616.11 1754.27,-2606.16 1750.9,-2616.2 1757.9,-2616.11\"/>\n",
+       "</g>\n",
+       "<!-- 68 -->\n",
+       "<g id=\"node69\" class=\"node\">\n",
+       "<title>68</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1635.5,-2684 1486.5,-2684 1486.5,-2642 1635.5,-2642 1635.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1487,-2642 1487,-2684 1537,-2684 1537,-2642 1487,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1492\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1493.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1537,-2663 1537,-2684 1585,-2684 1585,-2663 1537,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1547\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1585,-2663 1585,-2684 1636,-2684 1636,-2663 1585,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1593\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 12) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1537,-2642 1537,-2663 1585,-2663 1585,-2642 1537,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1542\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1585,-2642 1585,-2663 1636,-2663 1636,-2642 1585,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1590\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 67&#45;&gt;68 -->\n",
+       "<g id=\"edge72\" class=\"edge\">\n",
+       "<title>67&#45;&gt;68</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1562.46,-2719.63C1562.26,-2711.82 1562.02,-2702.73 1561.79,-2694.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1565.29,-2694.06 1561.53,-2684.16 1558.29,-2694.25 1565.29,-2694.06\"/>\n",
+       "</g>\n",
+       "<!-- 81 -->\n",
+       "<g id=\"node82\" class=\"node\">\n",
+       "<title>81</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1633,-2606 1487,-2606 1487,-2564 1633,-2564 1633,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1487,-2564 1487,-2606 1534,-2606 1534,-2564 1487,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1495\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1492\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1534,-2585 1534,-2606 1582,-2606 1582,-2585 1534,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1544\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1582,-2585 1582,-2606 1633,-2606 1633,-2585 1582,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1587\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1534,-2564 1534,-2585 1582,-2585 1582,-2564 1534,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1539\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1582,-2564 1582,-2585 1633,-2585 1633,-2564 1582,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1587\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 68&#45;&gt;81 -->\n",
+       "<g id=\"edge87\" class=\"edge\">\n",
+       "<title>68&#45;&gt;81</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1560.73,-2641.63C1560.63,-2633.82 1560.51,-2624.73 1560.4,-2616.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1563.9,-2616.11 1560.27,-2606.16 1556.9,-2616.2 1563.9,-2616.11\"/>\n",
+       "</g>\n",
+       "<!-- 70 -->\n",
+       "<g id=\"node71\" class=\"node\">\n",
+       "<title>70</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1121.5,-2684 978.5,-2684 978.5,-2642 1121.5,-2642 1121.5,-2684\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"979,-2642 979,-2684 1029,-2684 1029,-2642 979,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"984\" y=\"-2666\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"985.5\" y=\"-2655\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1029,-2663 1029,-2684 1077,-2684 1077,-2663 1029,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1039\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1077,-2663 1077,-2684 1122,-2684 1122,-2663 1077,-2663\"/>\n",
+       "<text text-anchor=\"start\" x=\"1082\" y=\"-2671\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1029,-2642 1029,-2663 1077,-2663 1077,-2642 1029,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1034\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1077,-2642 1077,-2663 1122,-2663 1122,-2642 1077,-2642\"/>\n",
+       "<text text-anchor=\"start\" x=\"1082\" y=\"-2650\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 69&#45;&gt;70 -->\n",
+       "<g id=\"edge74\" class=\"edge\">\n",
+       "<title>69&#45;&gt;70</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1014.13,-2719.63C1019.45,-2711.39 1025.68,-2701.72 1031.44,-2692.78\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1034.53,-2694.46 1037,-2684.16 1028.64,-2690.67 1034.53,-2694.46\"/>\n",
+       "</g>\n",
+       "<!-- 71 -->\n",
+       "<g id=\"node72\" class=\"node\">\n",
+       "<title>71</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1177,-2606 1019,-2606 1019,-2564 1177,-2564 1177,-2606\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1019,-2564 1019,-2606 1084,-2606 1084,-2564 1019,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1024\" y=\"-2588\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1033\" y=\"-2577\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1084,-2585 1084,-2606 1132,-2606 1132,-2585 1084,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1094\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1132,-2585 1132,-2606 1177,-2606 1177,-2585 1132,-2585\"/>\n",
+       "<text text-anchor=\"start\" x=\"1137\" y=\"-2593\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1084,-2564 1084,-2585 1132,-2585 1132,-2564 1084,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1089\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1132,-2564 1132,-2585 1177,-2585 1177,-2564 1132,-2564\"/>\n",
+       "<text text-anchor=\"start\" x=\"1137\" y=\"-2572\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 70&#45;&gt;71 -->\n",
+       "<g id=\"edge75\" class=\"edge\">\n",
+       "<title>70&#45;&gt;71</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1062.86,-2641.63C1068.07,-2633.39 1074.18,-2623.72 1079.82,-2614.78\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1082.89,-2616.48 1085.27,-2606.16 1076.97,-2612.75 1082.89,-2616.48\"/>\n",
+       "</g>\n",
+       "<!-- 72 -->\n",
+       "<g id=\"node73\" class=\"node\">\n",
+       "<title>72</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1219,-1982 1079,-1982 1079,-1940 1219,-1940 1219,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1079,-1940 1079,-1982 1126,-1982 1126,-1940 1079,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1091\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1084\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1126,-1961 1126,-1982 1174,-1982 1174,-1961 1126,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1136\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1174,-1961 1174,-1982 1219,-1982 1219,-1961 1174,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1179\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1126,-1940 1126,-1961 1174,-1961 1174,-1940 1126,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1131\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1174,-1940 1174,-1961 1219,-1961 1219,-1940 1174,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1179\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
+       "</g>\n",
+       "<!-- 71&#45;&gt;72 -->\n",
+       "<g id=\"edge76\" class=\"edge\">\n",
+       "<title>71&#45;&gt;72</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1103.42,-2563.74C1110.78,-2534.46 1123,-2478.48 1123,-2430 1123,-2430 1123,-2430 1123,-2116 1123,-2072.39 1133.26,-2022.8 1140.89,-1991.95\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1144.29,-1992.8 1143.36,-1982.24 1137.5,-1991.07 1144.29,-1992.8\"/>\n",
+       "</g>\n",
+       "<!-- 102 -->\n",
+       "<g id=\"node103\" class=\"node\">\n",
+       "<title>102</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1483.5,-1748 1256.5,-1748 1256.5,-1706 1483.5,-1706 1483.5,-1748\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1257,-1706 1257,-1748 1304,-1748 1304,-1706 1257,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1272.5\" y=\"-1730\" font-family=\"Linux libertine\" font-size=\"10.00\">cat</text>\n",
+       "<text text-anchor=\"start\" x=\"1262\" y=\"-1719\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1304,-1727 1304,-1748 1352,-1748 1352,-1727 1304,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1314\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1352,-1727 1352,-1748 1484,-1748 1484,-1727 1352,-1727\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-1735\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10), 2 x (1, 32), (1, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1304,-1706 1304,-1727 1352,-1727 1352,-1706 1304,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1309\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1352,-1706 1352,-1727 1484,-1727 1484,-1706 1352,-1706\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400.5\" y=\"-1714\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "</g>\n",
+       "<!-- 72&#45;&gt;102 -->\n",
+       "<g id=\"edge77\" class=\"edge\">\n",
+       "<title>72&#45;&gt;102</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1160.95,-1939.99C1182.31,-1905.33 1229.83,-1833.13 1283,-1784 1295.34,-1772.59 1310.29,-1762.04 1324.28,-1753.25\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1326.14,-1756.21 1332.84,-1748 1322.49,-1750.24 1326.14,-1756.21\"/>\n",
+       "</g>\n",
+       "<!-- 74 -->\n",
+       "<g id=\"node75\" class=\"node\">\n",
+       "<title>74</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1450.5,-2528 1289.5,-2528 1289.5,-2486 1450.5,-2486 1450.5,-2528\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1290,-2486 1290,-2528 1340,-2528 1340,-2486 1290,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1295\" y=\"-2510\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1296.5\" y=\"-2499\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2507 1340,-2528 1388,-2528 1388,-2507 1340,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"1350\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2507 1388,-2528 1451,-2528 1451,-2507 1388,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2486 1340,-2507 1388,-2507 1388,-2486 1340,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1345\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2486 1388,-2507 1451,-2507 1451,-2486 1388,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 73&#45;&gt;74 -->\n",
+       "<g id=\"edge79\" class=\"edge\">\n",
+       "<title>73&#45;&gt;74</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2563.63C1370,-2555.82 1370,-2546.73 1370,-2538.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2538.16 1370,-2528.16 1366.5,-2538.16 1373.5,-2538.16\"/>\n",
+       "</g>\n",
+       "<!-- 75 -->\n",
+       "<g id=\"node76\" class=\"node\">\n",
+       "<title>75</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1458,-2450 1282,-2450 1282,-2408 1458,-2408 1458,-2450\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1282,-2408 1282,-2450 1347,-2450 1347,-2408 1282,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1287\" y=\"-2432\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2421\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-2429 1347,-2450 1395,-2450 1395,-2429 1347,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-2429 1395,-2450 1458,-2450 1458,-2429 1395,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-2408 1347,-2429 1395,-2429 1395,-2408 1347,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1352\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-2408 1395,-2429 1458,-2429 1458,-2408 1395,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 74&#45;&gt;75 -->\n",
+       "<g id=\"edge80\" class=\"edge\">\n",
+       "<title>74&#45;&gt;75</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2485.63C1370,-2477.82 1370,-2468.73 1370,-2460.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2460.16 1370,-2450.16 1366.5,-2460.16 1373.5,-2460.16\"/>\n",
+       "</g>\n",
+       "<!-- 76 -->\n",
+       "<g id=\"node77\" class=\"node\">\n",
+       "<title>76</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-2372 1291,-2372 1291,-2330 1449,-2330 1449,-2372\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-2330 1291,-2372 1338,-2372 1338,-2330 1291,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1303\" y=\"-2354\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2343\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2351 1338,-2372 1386,-2372 1386,-2351 1338,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2351 1386,-2372 1449,-2372 1449,-2351 1386,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2330 1338,-2351 1386,-2351 1386,-2330 1338,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1343\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2330 1386,-2351 1449,-2351 1449,-2330 1386,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "</g>\n",
+       "<!-- 75&#45;&gt;76 -->\n",
+       "<g id=\"edge81\" class=\"edge\">\n",
+       "<title>75&#45;&gt;76</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2407.63C1370,-2399.82 1370,-2390.73 1370,-2382.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2382.16 1370,-2372.16 1366.5,-2382.16 1373.5,-2382.16\"/>\n",
+       "</g>\n",
+       "<!-- 77 -->\n",
+       "<g id=\"node78\" class=\"node\">\n",
+       "<title>77</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-2294 1291,-2294 1291,-2252 1449,-2252 1449,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-2252 1291,-2294 1338,-2294 1338,-2252 1291,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1299\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2273 1338,-2294 1386,-2294 1386,-2273 1338,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2273 1386,-2294 1449,-2294 1449,-2273 1386,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2252 1338,-2273 1386,-2273 1386,-2252 1338,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1343\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2252 1386,-2273 1449,-2273 1449,-2252 1386,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 76&#45;&gt;77 -->\n",
+       "<g id=\"edge82\" class=\"edge\">\n",
+       "<title>76&#45;&gt;77</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2329.63C1370,-2321.82 1370,-2312.73 1370,-2304.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2304.16 1370,-2294.16 1366.5,-2304.16 1373.5,-2304.16\"/>\n",
+       "</g>\n",
+       "<!-- 78 -->\n",
+       "<g id=\"node79\" class=\"node\">\n",
+       "<title>78</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1450.5,-2216 1289.5,-2216 1289.5,-2174 1450.5,-2174 1450.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1290,-2174 1290,-2216 1340,-2216 1340,-2174 1290,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1295\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1296.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2195 1340,-2216 1388,-2216 1388,-2195 1340,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1350\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2195 1388,-2216 1451,-2216 1451,-2195 1388,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1340,-2174 1340,-2195 1388,-2195 1388,-2174 1340,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1345\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-2174 1388,-2195 1451,-2195 1451,-2174 1388,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 77&#45;&gt;78 -->\n",
+       "<g id=\"edge83\" class=\"edge\">\n",
+       "<title>77&#45;&gt;78</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2251.63C1370,-2243.82 1370,-2234.73 1370,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2226.16 1370,-2216.16 1366.5,-2226.16 1373.5,-2226.16\"/>\n",
+       "</g>\n",
+       "<!-- 79 -->\n",
+       "<g id=\"node80\" class=\"node\">\n",
+       "<title>79</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1458,-2138 1282,-2138 1282,-2096 1458,-2096 1458,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1282,-2096 1282,-2138 1347,-2138 1347,-2096 1282,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1287\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-2117 1347,-2138 1395,-2138 1395,-2117 1347,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-2117 1395,-2138 1458,-2138 1458,-2117 1395,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-2096 1347,-2117 1395,-2117 1395,-2096 1347,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1352\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-2096 1395,-2117 1458,-2117 1458,-2096 1395,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 78&#45;&gt;79 -->\n",
+       "<g id=\"edge84\" class=\"edge\">\n",
+       "<title>78&#45;&gt;79</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2173.63C1370,-2165.82 1370,-2156.73 1370,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2148.16 1370,-2138.16 1366.5,-2148.16 1373.5,-2148.16\"/>\n",
+       "</g>\n",
+       "<!-- 80 -->\n",
+       "<g id=\"node81\" class=\"node\">\n",
+       "<title>80</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-2060 1291,-2060 1291,-2018 1449,-2018 1449,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-2018 1291,-2060 1338,-2060 1338,-2018 1291,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1303\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2039 1338,-2060 1386,-2060 1386,-2039 1338,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2039 1386,-2060 1449,-2060 1449,-2039 1386,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1338,-2018 1338,-2039 1386,-2039 1386,-2018 1338,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1343\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1386,-2018 1386,-2039 1449,-2039 1449,-2018 1386,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1391\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "</g>\n",
+       "<!-- 79&#45;&gt;80 -->\n",
+       "<g id=\"edge85\" class=\"edge\">\n",
+       "<title>79&#45;&gt;80</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2095.63C1370,-2087.82 1370,-2078.73 1370,-2070.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-2070.16 1370,-2060.16 1366.5,-2070.16 1373.5,-2070.16\"/>\n",
+       "</g>\n",
+       "<!-- 93 -->\n",
+       "<g id=\"node94\" class=\"node\">\n",
+       "<title>93</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1451.5,-1982 1288.5,-1982 1288.5,-1940 1451.5,-1940 1451.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1289,-1940 1289,-1982 1341,-1982 1341,-1940 1289,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1294\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"1296.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1341,-1961 1341,-1982 1389,-1982 1389,-1961 1341,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1351\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1389,-1961 1389,-1982 1452,-1982 1452,-1961 1389,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1394\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 16, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1341,-1940 1341,-1961 1389,-1961 1389,-1940 1341,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1346\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1389,-1940 1389,-1961 1452,-1961 1452,-1940 1389,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1394\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "</g>\n",
+       "<!-- 80&#45;&gt;93 -->\n",
+       "<g id=\"edge86\" class=\"edge\">\n",
+       "<title>80&#45;&gt;93</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-2017.63C1370,-2009.82 1370,-2000.73 1370,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1992.16 1370,-1982.16 1366.5,-1992.16 1373.5,-1992.16\"/>\n",
+       "</g>\n",
+       "<!-- 82 -->\n",
+       "<g id=\"node83\" class=\"node\">\n",
+       "<title>82</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1634.5,-2528 1485.5,-2528 1485.5,-2486 1634.5,-2486 1634.5,-2528\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1486,-2486 1486,-2528 1536,-2528 1536,-2486 1486,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1491\" y=\"-2510\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1492.5\" y=\"-2499\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1536,-2507 1536,-2528 1584,-2528 1584,-2507 1536,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"1546\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1584,-2507 1584,-2528 1635,-2528 1635,-2507 1584,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"1589\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1536,-2486 1536,-2507 1584,-2507 1584,-2486 1536,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1541\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1584,-2486 1584,-2507 1635,-2507 1635,-2486 1584,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1589\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 81&#45;&gt;82 -->\n",
+       "<g id=\"edge88\" class=\"edge\">\n",
+       "<title>81&#45;&gt;82</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1560,-2563.63C1560,-2555.82 1560,-2546.73 1560,-2538.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1563.5,-2538.16 1560,-2528.16 1556.5,-2538.16 1563.5,-2538.16\"/>\n",
+       "</g>\n",
+       "<!-- 83 -->\n",
+       "<g id=\"node84\" class=\"node\">\n",
+       "<title>83</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1641,-2450 1477,-2450 1477,-2408 1641,-2408 1641,-2450\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1477,-2408 1477,-2450 1542,-2450 1542,-2408 1477,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1482\" y=\"-2432\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1491\" y=\"-2421\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1542,-2429 1542,-2450 1590,-2450 1590,-2429 1542,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"1552\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1590,-2429 1590,-2450 1641,-2450 1641,-2429 1590,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"1595\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1542,-2408 1542,-2429 1590,-2429 1590,-2408 1542,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1547\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1590,-2408 1590,-2429 1641,-2429 1641,-2408 1590,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1595\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 82&#45;&gt;83 -->\n",
+       "<g id=\"edge89\" class=\"edge\">\n",
+       "<title>82&#45;&gt;83</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1559.73,-2485.63C1559.63,-2477.82 1559.51,-2468.73 1559.4,-2460.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1562.9,-2460.11 1559.27,-2450.16 1555.9,-2460.2 1562.9,-2460.11\"/>\n",
+       "</g>\n",
+       "<!-- 84 -->\n",
+       "<g id=\"node85\" class=\"node\">\n",
+       "<title>84</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1630,-2372 1484,-2372 1484,-2330 1630,-2330 1630,-2372\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1484,-2330 1484,-2372 1531,-2372 1531,-2330 1484,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1496\" y=\"-2354\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1489\" y=\"-2343\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1531,-2351 1531,-2372 1579,-2372 1579,-2351 1531,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"1541\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1579,-2351 1579,-2372 1630,-2372 1630,-2351 1579,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"1584\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1531,-2330 1531,-2351 1579,-2351 1579,-2330 1531,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1536\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1579,-2330 1579,-2351 1630,-2351 1630,-2330 1579,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1584\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "</g>\n",
+       "<!-- 83&#45;&gt;84 -->\n",
+       "<g id=\"edge90\" class=\"edge\">\n",
+       "<title>83&#45;&gt;84</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1558.46,-2407.63C1558.26,-2399.82 1558.02,-2390.73 1557.79,-2382.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1561.29,-2382.06 1557.53,-2372.16 1554.29,-2382.25 1561.29,-2382.06\"/>\n",
+       "</g>\n",
+       "<!-- 96 -->\n",
+       "<g id=\"node97\" class=\"node\">\n",
+       "<title>96</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1632.5,-2294 1481.5,-2294 1481.5,-2252 1632.5,-2252 1632.5,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1482,-2252 1482,-2294 1534,-2294 1534,-2252 1482,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1487\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"1489.5\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1534,-2273 1534,-2294 1582,-2294 1582,-2273 1534,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1544\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1582,-2273 1582,-2294 1633,-2294 1633,-2273 1582,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1587\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 4, 3) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1534,-2252 1534,-2273 1582,-2273 1582,-2252 1534,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1539\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1582,-2252 1582,-2273 1633,-2273 1633,-2252 1582,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1587\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "</g>\n",
+       "<!-- 84&#45;&gt;96 -->\n",
+       "<g id=\"edge91\" class=\"edge\">\n",
+       "<title>84&#45;&gt;96</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1557,-2329.63C1557,-2321.82 1557,-2312.73 1557,-2304.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1560.5,-2304.16 1557,-2294.16 1553.5,-2304.16 1560.5,-2304.16\"/>\n",
+       "</g>\n",
+       "<!-- 86 -->\n",
+       "<g id=\"node87\" class=\"node\">\n",
+       "<title>86</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1835.5,-2528 1668.5,-2528 1668.5,-2486 1835.5,-2486 1835.5,-2528\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1669,-2486 1669,-2528 1719,-2528 1719,-2486 1669,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1674\" y=\"-2510\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1675.5\" y=\"-2499\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1719,-2507 1719,-2528 1767,-2528 1767,-2507 1719,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"1729\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1767,-2507 1767,-2528 1836,-2528 1836,-2507 1767,-2507\"/>\n",
+       "<text text-anchor=\"start\" x=\"1772\" y=\"-2515\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1719,-2486 1719,-2507 1767,-2507 1767,-2486 1719,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1724\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1767,-2486 1767,-2507 1836,-2507 1836,-2486 1767,-2486\"/>\n",
+       "<text text-anchor=\"start\" x=\"1772\" y=\"-2494\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 85&#45;&gt;86 -->\n",
+       "<g id=\"edge93\" class=\"edge\">\n",
+       "<title>85&#45;&gt;86</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1753.46,-2563.63C1753.26,-2555.82 1753.02,-2546.73 1752.79,-2538.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1756.29,-2538.06 1752.53,-2528.16 1749.29,-2538.25 1756.29,-2538.06\"/>\n",
+       "</g>\n",
+       "<!-- 87 -->\n",
+       "<g id=\"node88\" class=\"node\">\n",
+       "<title>87</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1843,-2450 1661,-2450 1661,-2408 1843,-2408 1843,-2450\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1661,-2408 1661,-2450 1726,-2450 1726,-2408 1661,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1666\" y=\"-2432\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1675\" y=\"-2421\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1726,-2429 1726,-2450 1774,-2450 1774,-2429 1726,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"1736\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1774,-2429 1774,-2450 1843,-2450 1843,-2429 1774,-2429\"/>\n",
+       "<text text-anchor=\"start\" x=\"1779\" y=\"-2437\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1726,-2408 1726,-2429 1774,-2429 1774,-2408 1726,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1731\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1774,-2408 1774,-2429 1843,-2429 1843,-2408 1774,-2408\"/>\n",
+       "<text text-anchor=\"start\" x=\"1779\" y=\"-2416\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 86&#45;&gt;87 -->\n",
+       "<g id=\"edge94\" class=\"edge\">\n",
+       "<title>86&#45;&gt;87</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1752,-2485.63C1752,-2477.82 1752,-2468.73 1752,-2460.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1755.5,-2460.16 1752,-2450.16 1748.5,-2460.16 1755.5,-2460.16\"/>\n",
+       "</g>\n",
+       "<!-- 88 -->\n",
+       "<g id=\"node89\" class=\"node\">\n",
+       "<title>88</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1833,-2372 1669,-2372 1669,-2330 1833,-2330 1833,-2372\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1669,-2330 1669,-2372 1716,-2372 1716,-2330 1669,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1681\" y=\"-2354\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1674\" y=\"-2343\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1716,-2351 1716,-2372 1764,-2372 1764,-2351 1716,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"1726\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1764,-2351 1764,-2372 1833,-2372 1833,-2351 1764,-2351\"/>\n",
+       "<text text-anchor=\"start\" x=\"1769\" y=\"-2359\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1716,-2330 1716,-2351 1764,-2351 1764,-2330 1716,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1721\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1764,-2330 1764,-2351 1833,-2351 1833,-2330 1764,-2330\"/>\n",
+       "<text text-anchor=\"start\" x=\"1769\" y=\"-2338\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "</g>\n",
+       "<!-- 87&#45;&gt;88 -->\n",
+       "<g id=\"edge95\" class=\"edge\">\n",
+       "<title>87&#45;&gt;88</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1751.73,-2407.63C1751.63,-2399.82 1751.51,-2390.73 1751.4,-2382.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1754.9,-2382.11 1751.27,-2372.16 1747.9,-2382.2 1754.9,-2382.11\"/>\n",
+       "</g>\n",
+       "<!-- 89 -->\n",
+       "<g id=\"node90\" class=\"node\">\n",
+       "<title>89</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1831,-2294 1667,-2294 1667,-2252 1831,-2252 1831,-2294\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1667,-2252 1667,-2294 1714,-2294 1714,-2252 1667,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1675\" y=\"-2276\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1672\" y=\"-2265\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1714,-2273 1714,-2294 1762,-2294 1762,-2273 1714,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1724\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1762,-2273 1762,-2294 1831,-2294 1831,-2273 1762,-2273\"/>\n",
+       "<text text-anchor=\"start\" x=\"1767\" y=\"-2281\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1714,-2252 1714,-2273 1762,-2273 1762,-2252 1714,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1719\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1762,-2252 1762,-2273 1831,-2273 1831,-2252 1762,-2252\"/>\n",
+       "<text text-anchor=\"start\" x=\"1767\" y=\"-2260\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 88&#45;&gt;89 -->\n",
+       "<g id=\"edge96\" class=\"edge\">\n",
+       "<title>88&#45;&gt;89</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1750.46,-2329.63C1750.26,-2321.82 1750.02,-2312.73 1749.79,-2304.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1753.29,-2304.06 1749.53,-2294.16 1746.29,-2304.25 1753.29,-2304.06\"/>\n",
+       "</g>\n",
+       "<!-- 90 -->\n",
+       "<g id=\"node91\" class=\"node\">\n",
+       "<title>90</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1831.5,-2216 1664.5,-2216 1664.5,-2174 1831.5,-2174 1831.5,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1665,-2174 1665,-2216 1715,-2216 1715,-2174 1665,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1670\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1671.5\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1715,-2195 1715,-2216 1763,-2216 1763,-2195 1715,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1725\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1763,-2195 1763,-2216 1832,-2216 1832,-2195 1763,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1768\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1715,-2174 1715,-2195 1763,-2195 1763,-2174 1715,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1720\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1763,-2174 1763,-2195 1832,-2195 1832,-2174 1763,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1768\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 89&#45;&gt;90 -->\n",
+       "<g id=\"edge97\" class=\"edge\">\n",
+       "<title>89&#45;&gt;90</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1748.73,-2251.63C1748.63,-2243.82 1748.51,-2234.73 1748.4,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1751.9,-2226.11 1748.27,-2216.16 1744.9,-2226.2 1751.9,-2226.11\"/>\n",
+       "</g>\n",
+       "<!-- 91 -->\n",
+       "<g id=\"node92\" class=\"node\">\n",
+       "<title>91</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1838,-2138 1656,-2138 1656,-2096 1838,-2096 1838,-2138\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1656,-2096 1656,-2138 1721,-2138 1721,-2096 1656,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1661\" y=\"-2120\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1670\" y=\"-2109\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1721,-2117 1721,-2138 1769,-2138 1769,-2117 1721,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1731\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1769,-2117 1769,-2138 1838,-2138 1838,-2117 1769,-2117\"/>\n",
+       "<text text-anchor=\"start\" x=\"1774\" y=\"-2125\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1721,-2096 1721,-2117 1769,-2117 1769,-2096 1721,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1726\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1769,-2096 1769,-2117 1838,-2117 1838,-2096 1769,-2096\"/>\n",
+       "<text text-anchor=\"start\" x=\"1774\" y=\"-2104\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 90&#45;&gt;91 -->\n",
+       "<g id=\"edge98\" class=\"edge\">\n",
+       "<title>90&#45;&gt;91</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1747.73,-2173.63C1747.63,-2165.82 1747.51,-2156.73 1747.4,-2148.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1750.9,-2148.11 1747.27,-2138.16 1743.9,-2148.2 1750.9,-2148.11\"/>\n",
+       "</g>\n",
+       "<!-- 92 -->\n",
+       "<g id=\"node93\" class=\"node\">\n",
+       "<title>92</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1811,-2060 1647,-2060 1647,-2018 1811,-2018 1811,-2060\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1647,-2018 1647,-2060 1694,-2060 1694,-2018 1647,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1659\" y=\"-2042\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1652\" y=\"-2031\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1694,-2039 1694,-2060 1742,-2060 1742,-2039 1694,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1704\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1742,-2039 1742,-2060 1811,-2060 1811,-2039 1742,-2039\"/>\n",
+       "<text text-anchor=\"start\" x=\"1747\" y=\"-2047\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1694,-2018 1694,-2039 1742,-2039 1742,-2018 1694,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1699\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1742,-2018 1742,-2039 1811,-2039 1811,-2018 1742,-2018\"/>\n",
+       "<text text-anchor=\"start\" x=\"1747\" y=\"-2026\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "</g>\n",
+       "<!-- 91&#45;&gt;92 -->\n",
+       "<g id=\"edge99\" class=\"edge\">\n",
+       "<title>91&#45;&gt;92</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1742.18,-2095.63C1740.31,-2087.73 1738.12,-2078.53 1736.08,-2069.9\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1739.49,-2069.08 1733.77,-2060.16 1732.67,-2070.7 1739.49,-2069.08\"/>\n",
+       "</g>\n",
+       "<!-- 99 -->\n",
+       "<g id=\"node100\" class=\"node\">\n",
+       "<title>99</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1811.5,-1982 1642.5,-1982 1642.5,-1940 1811.5,-1940 1811.5,-1982\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1643,-1940 1643,-1982 1695,-1982 1695,-1940 1643,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1648\" y=\"-1964\" font-family=\"Linux libertine\" font-size=\"10.00\">permute</text>\n",
+       "<text text-anchor=\"start\" x=\"1650.5\" y=\"-1953\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1695,-1961 1695,-1982 1743,-1982 1743,-1961 1695,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1705\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1743,-1961 1743,-1982 1812,-1982 1812,-1961 1743,-1961\"/>\n",
+       "<text text-anchor=\"start\" x=\"1748\" y=\"-1969\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 500, 32) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1695,-1940 1695,-1961 1743,-1961 1743,-1940 1695,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1700\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1743,-1940 1743,-1961 1812,-1961 1812,-1940 1743,-1940\"/>\n",
+       "<text text-anchor=\"start\" x=\"1748\" y=\"-1948\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "</g>\n",
+       "<!-- 92&#45;&gt;99 -->\n",
+       "<g id=\"edge100\" class=\"edge\">\n",
+       "<title>92&#45;&gt;99</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1728.46,-2017.63C1728.26,-2009.82 1728.02,-2000.73 1727.79,-1992.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1731.29,-1992.06 1727.53,-1982.16 1724.29,-1992.25 1731.29,-1992.06\"/>\n",
+       "</g>\n",
+       "<!-- 94 -->\n",
+       "<g id=\"node95\" class=\"node\">\n",
+       "<title>94</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1461,-1904 1279,-1904 1279,-1862 1461,-1862 1461,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1279,-1862 1279,-1904 1350,-1904 1350,-1862 1279,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1284\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1350,-1883 1350,-1904 1398,-1904 1398,-1883 1350,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1360\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-1883 1398,-1904 1461,-1904 1461,-1883 1398,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 16) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1350,-1862 1350,-1883 1398,-1883 1398,-1862 1350,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1355\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1398,-1862 1398,-1883 1461,-1883 1461,-1862 1398,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 93&#45;&gt;94 -->\n",
+       "<g id=\"edge101\" class=\"edge\">\n",
+       "<title>93&#45;&gt;94</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1939.63C1370,-1931.82 1370,-1922.73 1370,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1914.16 1370,-1904.16 1366.5,-1914.16 1373.5,-1914.16\"/>\n",
+       "</g>\n",
+       "<!-- 95 -->\n",
+       "<g id=\"node96\" class=\"node\">\n",
+       "<title>95</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1448,-1826 1292,-1826 1292,-1784 1448,-1784 1448,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1292,-1784 1292,-1826 1343,-1826 1343,-1784 1292,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1297\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1299\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1343,-1805 1343,-1826 1391,-1826 1391,-1805 1343,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1353\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1391,-1805 1391,-1826 1448,-1826 1448,-1805 1391,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1396\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1343,-1784 1343,-1805 1391,-1805 1391,-1784 1343,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1348\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1391,-1784 1391,-1805 1448,-1805 1448,-1784 1391,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1402\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 94&#45;&gt;95 -->\n",
+       "<g id=\"edge102\" class=\"edge\">\n",
+       "<title>94&#45;&gt;95</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1861.63C1370,-1853.82 1370,-1844.73 1370,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1836.16 1370,-1826.16 1366.5,-1836.16 1373.5,-1836.16\"/>\n",
+       "</g>\n",
+       "<!-- 95&#45;&gt;102 -->\n",
+       "<g id=\"edge103\" class=\"edge\">\n",
+       "<title>95&#45;&gt;102</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1783.63C1370,-1775.82 1370,-1766.73 1370,-1758.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1758.16 1370,-1748.16 1366.5,-1758.16 1373.5,-1758.16\"/>\n",
+       "</g>\n",
+       "<!-- 97 -->\n",
+       "<g id=\"node98\" class=\"node\">\n",
+       "<title>97</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1642,-2216 1472,-2216 1472,-2174 1642,-2174 1642,-2216\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1472,-2174 1472,-2216 1543,-2216 1543,-2174 1472,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1477\" y=\"-2198\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"1489\" y=\"-2187\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1543,-2195 1543,-2216 1591,-2216 1591,-2195 1543,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1553\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1591,-2195 1591,-2216 1642,-2216 1642,-2195 1591,-2195\"/>\n",
+       "<text text-anchor=\"start\" x=\"1596\" y=\"-2203\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 4) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1543,-2174 1543,-2195 1591,-2195 1591,-2174 1543,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1548\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1591,-2174 1591,-2195 1642,-2195 1642,-2174 1591,-2174\"/>\n",
+       "<text text-anchor=\"start\" x=\"1596\" y=\"-2182\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "</g>\n",
+       "<!-- 96&#45;&gt;97 -->\n",
+       "<g id=\"edge104\" class=\"edge\">\n",
+       "<title>96&#45;&gt;97</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1557,-2251.63C1557,-2243.82 1557,-2234.73 1557,-2226.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1560.5,-2226.16 1557,-2216.16 1553.5,-2226.16 1560.5,-2226.16\"/>\n",
+       "</g>\n",
+       "<!-- 98 -->\n",
+       "<g id=\"node99\" class=\"node\">\n",
+       "<title>98</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1620,-1826 1470,-1826 1470,-1784 1620,-1784 1620,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1470,-1784 1470,-1826 1521,-1826 1521,-1784 1470,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1475\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1477\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1521,-1805 1521,-1826 1569,-1826 1569,-1805 1521,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1531\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1569,-1805 1569,-1826 1620,-1826 1620,-1805 1569,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1574\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1521,-1784 1521,-1805 1569,-1805 1569,-1784 1521,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1526\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1569,-1784 1569,-1805 1620,-1805 1620,-1784 1569,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1580\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 3) </text>\n",
+       "</g>\n",
+       "<!-- 97&#45;&gt;98 -->\n",
+       "<g id=\"edge105\" class=\"edge\">\n",
+       "<title>97&#45;&gt;98</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1555.26,-2173.59C1552.91,-2144.12 1549,-2087.92 1549,-2040 1549,-2040 1549,-2040 1549,-1960 1549,-1917.18 1547.44,-1867.67 1546.26,-1836.6\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1549.75,-1836.28 1545.87,-1826.42 1542.76,-1836.55 1549.75,-1836.28\"/>\n",
+       "</g>\n",
+       "<!-- 98&#45;&gt;102 -->\n",
+       "<g id=\"edge106\" class=\"edge\">\n",
+       "<title>98&#45;&gt;102</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1498.56,-1783.83C1476.22,-1774.13 1449.27,-1762.43 1425.82,-1752.24\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1426.92,-1748.9 1416.35,-1748.13 1424.13,-1755.32 1426.92,-1748.9\"/>\n",
+       "</g>\n",
+       "<!-- 100 -->\n",
+       "<g id=\"node101\" class=\"node\">\n",
+       "<title>100</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1820,-1904 1632,-1904 1632,-1862 1820,-1862 1820,-1904\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1632,-1862 1632,-1904 1703,-1904 1703,-1862 1632,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1637\" y=\"-1886\" font-family=\"Linux libertine\" font-size=\"10.00\">max_pool1d</text>\n",
+       "<text text-anchor=\"start\" x=\"1649\" y=\"-1875\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1703,-1883 1703,-1904 1751,-1904 1751,-1883 1703,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1713\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1751,-1883 1751,-1904 1820,-1904 1820,-1883 1751,-1883\"/>\n",
+       "<text text-anchor=\"start\" x=\"1756\" y=\"-1891\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 500) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1703,-1862 1703,-1883 1751,-1883 1751,-1862 1703,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1708\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1751,-1862 1751,-1883 1820,-1883 1820,-1862 1751,-1862\"/>\n",
+       "<text text-anchor=\"start\" x=\"1762\" y=\"-1870\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "</g>\n",
+       "<!-- 99&#45;&gt;100 -->\n",
+       "<g id=\"edge107\" class=\"edge\">\n",
+       "<title>99&#45;&gt;100</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1726.73,-1939.63C1726.63,-1931.82 1726.51,-1922.73 1726.4,-1914.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1729.9,-1914.11 1726.27,-1904.16 1722.9,-1914.2 1729.9,-1914.11\"/>\n",
+       "</g>\n",
+       "<!-- 101 -->\n",
+       "<g id=\"node102\" class=\"node\">\n",
+       "<title>101</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1801,-1826 1645,-1826 1645,-1784 1801,-1784 1801,-1826\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1645,-1784 1645,-1826 1696,-1826 1696,-1784 1645,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1650\" y=\"-1808\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1652\" y=\"-1797\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:2</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1696,-1805 1696,-1826 1744,-1826 1744,-1805 1696,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1706\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1744,-1805 1744,-1826 1801,-1826 1801,-1805 1744,-1805\"/>\n",
+       "<text text-anchor=\"start\" x=\"1749\" y=\"-1813\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1696,-1784 1696,-1805 1744,-1805 1744,-1784 1696,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1701\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1744,-1784 1744,-1805 1801,-1805 1801,-1784 1744,-1784\"/>\n",
+       "<text text-anchor=\"start\" x=\"1755\" y=\"-1792\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 32) </text>\n",
+       "</g>\n",
+       "<!-- 100&#45;&gt;101 -->\n",
+       "<g id=\"edge108\" class=\"edge\">\n",
+       "<title>100&#45;&gt;101</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1725.2,-1861.63C1724.89,-1853.82 1724.53,-1844.73 1724.19,-1836.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1727.69,-1836.01 1723.8,-1826.16 1720.69,-1836.29 1727.69,-1836.01\"/>\n",
+       "</g>\n",
+       "<!-- 101&#45;&gt;102 -->\n",
+       "<g id=\"edge109\" class=\"edge\">\n",
+       "<title>101&#45;&gt;102</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1644.66,-1787.13C1594.18,-1776.26 1527.99,-1762.01 1473.23,-1750.23\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1473.66,-1746.74 1463.15,-1748.06 1472.19,-1753.58 1473.66,-1746.74\"/>\n",
+       "</g>\n",
+       "<!-- 103 -->\n",
+       "<g id=\"node104\" class=\"node\">\n",
+       "<title>103</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1443,-1670 1297,-1670 1297,-1628 1443,-1628 1443,-1670\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1297,-1628 1297,-1670 1344,-1670 1344,-1628 1297,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-1652\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1641\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1649 1344,-1670 1392,-1670 1392,-1649 1344,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1354\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1649 1392,-1670 1443,-1670 1443,-1649 1392,-1649\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-1657\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 77) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1628 1344,-1649 1392,-1649 1392,-1628 1344,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1349\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1628 1392,-1649 1443,-1649 1443,-1628 1392,-1628\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1636\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 102&#45;&gt;103 -->\n",
+       "<g id=\"edge110\" class=\"edge\">\n",
+       "<title>102&#45;&gt;103</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1705.63C1370,-1697.82 1370,-1688.73 1370,-1680.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1680.16 1370,-1670.16 1366.5,-1680.16 1373.5,-1680.16\"/>\n",
+       "</g>\n",
+       "<!-- 104 -->\n",
+       "<g id=\"node105\" class=\"node\">\n",
+       "<title>104</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1452,-1592 1288,-1592 1288,-1550 1452,-1550 1452,-1592\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1288,-1550 1288,-1592 1353,-1592 1353,-1550 1288,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1293\" y=\"-1574\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1563\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-1571 1353,-1592 1401,-1592 1401,-1571 1353,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-1571 1401,-1592 1452,-1592 1452,-1571 1401,-1571\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-1579\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-1550 1353,-1571 1401,-1571 1401,-1550 1353,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-1550 1401,-1571 1452,-1571 1452,-1550 1401,-1550\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-1558\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 103&#45;&gt;104 -->\n",
+       "<g id=\"edge111\" class=\"edge\">\n",
+       "<title>103&#45;&gt;104</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1627.63C1370,-1619.82 1370,-1610.73 1370,-1602.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1602.16 1370,-1592.16 1366.5,-1602.16 1373.5,-1602.16\"/>\n",
+       "</g>\n",
+       "<!-- 105 -->\n",
+       "<g id=\"node106\" class=\"node\">\n",
+       "<title>105</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1443,-1514 1297,-1514 1297,-1472 1443,-1472 1443,-1514\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1297,-1472 1297,-1514 1344,-1514 1344,-1472 1297,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1309\" y=\"-1496\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1485\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1493 1344,-1514 1392,-1514 1392,-1493 1344,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"1354\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1493 1392,-1514 1443,-1514 1443,-1493 1392,-1493\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1501\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1472 1344,-1493 1392,-1493 1392,-1472 1344,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1349\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1472 1392,-1493 1443,-1493 1443,-1472 1392,-1472\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1480\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 104&#45;&gt;105 -->\n",
+       "<g id=\"edge112\" class=\"edge\">\n",
+       "<title>104&#45;&gt;105</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1549.63C1370,-1541.82 1370,-1532.73 1370,-1524.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1524.16 1370,-1514.16 1366.5,-1524.16 1373.5,-1524.16\"/>\n",
+       "</g>\n",
+       "<!-- 106 -->\n",
+       "<g id=\"node107\" class=\"node\">\n",
+       "<title>106</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1444.5,-1436 1295.5,-1436 1295.5,-1394 1444.5,-1394 1444.5,-1436\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1296,-1394 1296,-1436 1346,-1436 1346,-1394 1296,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1301\" y=\"-1418\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1302.5\" y=\"-1407\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1346,-1415 1346,-1436 1394,-1436 1394,-1415 1346,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1356\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1394,-1415 1394,-1436 1445,-1436 1445,-1415 1394,-1415\"/>\n",
+       "<text text-anchor=\"start\" x=\"1399\" y=\"-1423\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1346,-1394 1346,-1415 1394,-1415 1394,-1394 1346,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1351\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1394,-1394 1394,-1415 1445,-1415 1445,-1394 1394,-1394\"/>\n",
+       "<text text-anchor=\"start\" x=\"1399\" y=\"-1402\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "</g>\n",
+       "<!-- 105&#45;&gt;106 -->\n",
+       "<g id=\"edge113\" class=\"edge\">\n",
+       "<title>105&#45;&gt;106</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1471.63C1370,-1463.82 1370,-1454.73 1370,-1446.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1446.16 1370,-1436.16 1366.5,-1446.16 1373.5,-1446.16\"/>\n",
+       "</g>\n",
+       "<!-- 107 -->\n",
+       "<g id=\"node108\" class=\"node\">\n",
+       "<title>107</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1443,-1358 1297,-1358 1297,-1316 1443,-1316 1443,-1358\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1297,-1316 1297,-1358 1344,-1358 1344,-1316 1297,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-1340\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1329\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1337 1344,-1358 1392,-1358 1392,-1337 1344,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1354\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1337 1392,-1358 1443,-1358 1443,-1337 1392,-1337\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1345\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 256) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1316 1344,-1337 1392,-1337 1392,-1316 1344,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1349\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1316 1392,-1337 1443,-1337 1443,-1316 1392,-1316\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1324\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 106&#45;&gt;107 -->\n",
+       "<g id=\"edge114\" class=\"edge\">\n",
+       "<title>106&#45;&gt;107</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1393.63C1370,-1385.82 1370,-1376.73 1370,-1368.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1368.16 1370,-1358.16 1366.5,-1368.16 1373.5,-1368.16\"/>\n",
+       "</g>\n",
+       "<!-- 108 -->\n",
+       "<g id=\"node109\" class=\"node\">\n",
+       "<title>108</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1452,-1280 1288,-1280 1288,-1238 1452,-1238 1452,-1280\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1288,-1238 1288,-1280 1353,-1280 1353,-1238 1288,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1293\" y=\"-1262\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1251\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-1259 1353,-1280 1401,-1280 1401,-1259 1353,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1363\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-1259 1401,-1280 1452,-1280 1452,-1259 1401,-1259\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-1267\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-1238 1353,-1259 1401,-1259 1401,-1238 1353,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1401,-1238 1401,-1259 1452,-1259 1452,-1238 1401,-1238\"/>\n",
+       "<text text-anchor=\"start\" x=\"1406\" y=\"-1246\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 107&#45;&gt;108 -->\n",
+       "<g id=\"edge115\" class=\"edge\">\n",
+       "<title>107&#45;&gt;108</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1315.63C1370,-1307.82 1370,-1298.73 1370,-1290.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1290.16 1370,-1280.16 1366.5,-1290.16 1373.5,-1290.16\"/>\n",
+       "</g>\n",
+       "<!-- 109 -->\n",
+       "<g id=\"node110\" class=\"node\">\n",
+       "<title>109</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1443,-1202 1297,-1202 1297,-1160 1443,-1160 1443,-1202\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1297,-1160 1297,-1202 1344,-1202 1344,-1160 1297,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1309\" y=\"-1184\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1173\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1181 1344,-1202 1392,-1202 1392,-1181 1344,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"1354\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1181 1392,-1202 1443,-1202 1443,-1181 1392,-1181\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1189\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1160 1344,-1181 1392,-1181 1392,-1160 1344,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1349\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1160 1392,-1181 1443,-1181 1443,-1160 1392,-1160\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1168\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 108&#45;&gt;109 -->\n",
+       "<g id=\"edge116\" class=\"edge\">\n",
+       "<title>108&#45;&gt;109</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1237.63C1370,-1229.82 1370,-1220.73 1370,-1212.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1212.16 1370,-1202.16 1366.5,-1212.16 1373.5,-1212.16\"/>\n",
+       "</g>\n",
+       "<!-- 110 -->\n",
+       "<g id=\"node111\" class=\"node\">\n",
+       "<title>110</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1444.5,-1124 1295.5,-1124 1295.5,-1082 1444.5,-1082 1444.5,-1124\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1296,-1082 1296,-1124 1346,-1124 1346,-1082 1296,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1301\" y=\"-1106\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1302.5\" y=\"-1095\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1346,-1103 1346,-1124 1394,-1124 1394,-1103 1346,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"1356\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1394,-1103 1394,-1124 1445,-1124 1445,-1103 1394,-1103\"/>\n",
+       "<text text-anchor=\"start\" x=\"1399\" y=\"-1111\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1346,-1082 1346,-1103 1394,-1103 1394,-1082 1346,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1351\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1394,-1082 1394,-1103 1445,-1103 1445,-1082 1394,-1082\"/>\n",
+       "<text text-anchor=\"start\" x=\"1399\" y=\"-1090\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "</g>\n",
+       "<!-- 109&#45;&gt;110 -->\n",
+       "<g id=\"edge117\" class=\"edge\">\n",
+       "<title>109&#45;&gt;110</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1159.63C1370,-1151.82 1370,-1142.73 1370,-1134.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1134.16 1370,-1124.16 1366.5,-1134.16 1373.5,-1134.16\"/>\n",
+       "</g>\n",
+       "<!-- 111 -->\n",
+       "<g id=\"node112\" class=\"node\">\n",
+       "<title>111</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1443,-1046 1297,-1046 1297,-1004 1443,-1004 1443,-1046\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1297,-1004 1297,-1046 1344,-1046 1344,-1004 1297,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-1028\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1302\" y=\"-1017\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1025 1344,-1046 1392,-1046 1392,-1025 1344,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"1354\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1025 1392,-1046 1443,-1046 1443,-1025 1392,-1025\"/>\n",
+       "<text text-anchor=\"start\" x=\"1397\" y=\"-1033\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 128) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1344,-1004 1344,-1025 1392,-1025 1392,-1004 1344,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1349\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1392,-1004 1392,-1025 1443,-1025 1443,-1004 1392,-1004\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-1012\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 110&#45;&gt;111 -->\n",
+       "<g id=\"edge118\" class=\"edge\">\n",
+       "<title>110&#45;&gt;111</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1081.63C1370,-1073.82 1370,-1064.73 1370,-1056.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-1056.16 1370,-1046.16 1366.5,-1056.16 1373.5,-1056.16\"/>\n",
+       "</g>\n",
+       "<!-- 112 -->\n",
+       "<g id=\"node113\" class=\"node\">\n",
+       "<title>112</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-968 1291,-968 1291,-926 1449,-926 1449,-968\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-926 1291,-968 1356,-968 1356,-926 1291,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-950\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-939\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1356,-947 1356,-968 1404,-968 1404,-947 1356,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"1366\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1404,-947 1404,-968 1449,-968 1449,-947 1404,-947\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-955\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1356,-926 1356,-947 1404,-947 1404,-926 1356,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1361\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1404,-926 1404,-947 1449,-947 1449,-926 1404,-926\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-934\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 111&#45;&gt;112 -->\n",
+       "<g id=\"edge119\" class=\"edge\">\n",
+       "<title>111&#45;&gt;112</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-1003.63C1370,-995.82 1370,-986.73 1370,-978.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-978.16 1370,-968.16 1366.5,-978.16 1373.5,-978.16\"/>\n",
+       "</g>\n",
+       "<!-- 113 -->\n",
+       "<g id=\"node114\" class=\"node\">\n",
+       "<title>113</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1440,-890 1300,-890 1300,-848 1440,-848 1440,-890\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1300,-848 1300,-890 1347,-890 1347,-848 1300,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1312\" y=\"-872\" font-family=\"Linux libertine\" font-size=\"10.00\">Tanh</text>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-861\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-869 1347,-890 1395,-890 1395,-869 1347,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-869 1395,-890 1440,-890 1440,-869 1395,-869\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-877\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-848 1347,-869 1395,-869 1395,-848 1347,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1352\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-848 1395,-869 1440,-869 1440,-848 1395,-848\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-856\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 112&#45;&gt;113 -->\n",
+       "<g id=\"edge120\" class=\"edge\">\n",
+       "<title>112&#45;&gt;113</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-925.63C1370,-917.82 1370,-908.73 1370,-900.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-900.16 1370,-890.16 1366.5,-900.16 1373.5,-900.16\"/>\n",
+       "</g>\n",
+       "<!-- 114 -->\n",
+       "<g id=\"node115\" class=\"node\">\n",
+       "<title>114</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1441.5,-812 1298.5,-812 1298.5,-770 1441.5,-770 1441.5,-812\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1299,-770 1299,-812 1349,-812 1349,-770 1299,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1304\" y=\"-794\" font-family=\"Linux libertine\" font-size=\"10.00\">Dropout</text>\n",
+       "<text text-anchor=\"start\" x=\"1305.5\" y=\"-783\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1349,-791 1349,-812 1397,-812 1397,-791 1349,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"1359\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1397,-791 1397,-812 1442,-812 1442,-791 1397,-791\"/>\n",
+       "<text text-anchor=\"start\" x=\"1402\" y=\"-799\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1349,-770 1349,-791 1397,-791 1397,-770 1349,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1354\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1397,-770 1397,-791 1442,-791 1442,-770 1397,-770\"/>\n",
+       "<text text-anchor=\"start\" x=\"1402\" y=\"-778\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 113&#45;&gt;114 -->\n",
+       "<g id=\"edge121\" class=\"edge\">\n",
+       "<title>113&#45;&gt;114</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-847.63C1370,-839.82 1370,-830.73 1370,-822.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-822.16 1370,-812.16 1366.5,-822.16 1373.5,-822.16\"/>\n",
+       "</g>\n",
+       "<!-- 115 -->\n",
+       "<g id=\"node116\" class=\"node\">\n",
+       "<title>115</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1440,-734 1300,-734 1300,-692 1440,-692 1440,-734\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1300,-692 1300,-734 1347,-734 1347,-692 1300,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1308\" y=\"-716\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-705\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-713 1347,-734 1395,-734 1395,-713 1347,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-713 1395,-734 1440,-734 1440,-713 1395,-713\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-721\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-692 1347,-713 1395,-713 1395,-692 1347,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1352\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-692 1395,-713 1440,-713 1440,-692 1395,-692\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-700\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 114&#45;&gt;115 -->\n",
+       "<g id=\"edge122\" class=\"edge\">\n",
+       "<title>114&#45;&gt;115</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-769.63C1370,-761.82 1370,-752.73 1370,-744.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-744.16 1370,-734.16 1366.5,-744.16 1373.5,-744.16\"/>\n",
+       "</g>\n",
+       "<!-- 116 -->\n",
+       "<g id=\"node117\" class=\"node\">\n",
+       "<title>116</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1449,-656 1291,-656 1291,-614 1449,-614 1449,-656\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1291,-614 1291,-656 1356,-656 1356,-614 1291,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1296\" y=\"-638\" font-family=\"Linux libertine\" font-size=\"10.00\">LayerNorm</text>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-627\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:3</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1356,-635 1356,-656 1404,-656 1404,-635 1356,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1366\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1404,-635 1404,-656 1449,-656 1449,-635 1404,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1356,-614 1356,-635 1404,-635 1404,-614 1356,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1361\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1404,-614 1404,-635 1449,-635 1449,-614 1404,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1409\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "</g>\n",
+       "<!-- 115&#45;&gt;116 -->\n",
+       "<g id=\"edge123\" class=\"edge\">\n",
+       "<title>115&#45;&gt;116</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-691.63C1370,-683.82 1370,-674.73 1370,-666.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-666.16 1370,-656.16 1366.5,-666.16 1373.5,-666.16\"/>\n",
+       "</g>\n",
+       "<!-- 117 -->\n",
+       "<g id=\"node118\" class=\"node\">\n",
+       "<title>117</title>\n",
+       "<polygon fill=\"#c1ffc1\" stroke=\"transparent\" points=\"1440,-578 1300,-578 1300,-536 1440,-536 1440,-578\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1300,-536 1300,-578 1347,-578 1347,-536 1300,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1308\" y=\"-560\" font-family=\"Linux libertine\" font-size=\"10.00\">Linear</text>\n",
+       "<text text-anchor=\"start\" x=\"1305\" y=\"-549\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-557 1347,-578 1395,-578 1395,-557 1347,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1357\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-557 1395,-578 1440,-578 1440,-557 1395,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1400\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 64) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1347,-536 1347,-557 1395,-557 1395,-536 1347,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1352\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1395,-536 1395,-557 1440,-557 1440,-536 1395,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1403\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 116&#45;&gt;117 -->\n",
+       "<g id=\"edge124\" class=\"edge\">\n",
+       "<title>116&#45;&gt;117</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-613.63C1370,-605.82 1370,-596.73 1370,-588.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-588.16 1370,-578.16 1366.5,-588.16 1373.5,-588.16\"/>\n",
+       "</g>\n",
+       "<!-- 118 -->\n",
+       "<g id=\"node119\" class=\"node\">\n",
+       "<title>118</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"1426.5,-495 1313.5,-495 1313.5,-463 1426.5,-463 1426.5,-495\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1314,-463 1314,-495 1391,-495 1391,-463 1314,-463\"/>\n",
+       "<text text-anchor=\"start\" x=\"1319\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"1334\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1391,-463 1391,-495 1427,-495 1427,-463 1391,-463\"/>\n",
+       "<text text-anchor=\"start\" x=\"1396\" y=\"-476.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1)</text>\n",
+       "</g>\n",
+       "<!-- 117&#45;&gt;118 -->\n",
+       "<g id=\"edge125\" class=\"edge\">\n",
+       "<title>117&#45;&gt;118</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1370,-535.63C1370,-526.33 1370,-515.22 1370,-505.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1373.5,-505.36 1370,-495.36 1366.5,-505.36 1373.5,-505.36\"/>\n",
+       "</g>\n",
+       "<!-- 120 -->\n",
+       "<g id=\"node121\" class=\"node\">\n",
+       "<title>120</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1201,-500 1043,-500 1043,-458 1201,-458 1201,-500\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1043,-458 1043,-500 1108,-500 1108,-458 1043,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1048\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">logsumexp</text>\n",
+       "<text text-anchor=\"start\" x=\"1057\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1108,-479 1108,-500 1156,-500 1156,-479 1108,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1118\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1156,-479 1156,-500 1201,-500 1201,-479 1156,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1161\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1108,-458 1108,-479 1156,-479 1156,-458 1108,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1113\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1156,-458 1156,-479 1201,-479 1201,-458 1156,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1164\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 119&#45;&gt;120 -->\n",
+       "<g id=\"edge126\" class=\"edge\">\n",
+       "<title>119&#45;&gt;120</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1043.39,-535.83C1056.21,-526.73 1071.52,-515.86 1085.2,-506.14\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1087.35,-508.91 1093.47,-500.26 1083.29,-503.2 1087.35,-508.91\"/>\n",
+       "</g>\n",
+       "<!-- 121 -->\n",
+       "<g id=\"node122\" class=\"node\">\n",
+       "<title>121</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1208,-422 1036,-422 1036,-380 1208,-380 1208,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1036,-380 1036,-422 1083,-422 1083,-380 1036,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1050.5\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">sub</text>\n",
+       "<text text-anchor=\"start\" x=\"1041\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1083,-401 1083,-422 1131,-422 1131,-401 1083,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1093\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1131,-401 1131,-422 1208,-422 1208,-401 1131,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1136\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25), (1, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1083,-380 1083,-401 1131,-401 1131,-380 1083,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1088\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1131,-380 1131,-401 1208,-401 1208,-380 1131,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1152\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 119&#45;&gt;121 -->\n",
+       "<g id=\"edge127\" class=\"edge\">\n",
+       "<title>119&#45;&gt;121</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1014.41,-535.84C1014.75,-514.7 1018.1,-481.35 1034,-458 1042.36,-445.72 1054.41,-435.55 1066.91,-427.41\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1068.9,-430.3 1075.61,-422.1 1065.25,-424.33 1068.9,-430.3\"/>\n",
+       "</g>\n",
+       "<!-- 120&#45;&gt;121 -->\n",
+       "<g id=\"edge128\" class=\"edge\">\n",
+       "<title>120&#45;&gt;121</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1122,-457.63C1122,-449.82 1122,-440.73 1122,-432.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1125.5,-432.16 1122,-422.16 1118.5,-432.16 1125.5,-432.16\"/>\n",
+       "</g>\n",
+       "<!-- 122 -->\n",
+       "<g id=\"node123\" class=\"node\">\n",
+       "<title>122</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1032,-344 874,-344 874,-302 1032,-302 1032,-344\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"874,-302 874,-344 921,-344 921,-302 874,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"891\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">eq</text>\n",
+       "<text text-anchor=\"start\" x=\"879\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"921,-323 921,-344 969,-344 969,-323 921,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"931\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"969,-323 969,-344 1032,-344 1032,-323 969,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"974\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">2 x (1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"921,-302 921,-323 969,-323 969,-302 921,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"926\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"969,-302 969,-323 1032,-323 1032,-302 969,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"983\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 121&#45;&gt;122 -->\n",
+       "<g id=\"edge129\" class=\"edge\">\n",
+       "<title>121&#45;&gt;122</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1077.15,-379.83C1055.58,-370.13 1029.55,-358.43 1006.91,-348.24\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1008.32,-345.04 997.76,-344.13 1005.45,-351.42 1008.32,-345.04\"/>\n",
+       "</g>\n",
+       "<!-- 126 -->\n",
+       "<g id=\"node127\" class=\"node\">\n",
+       "<title>126</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1193.5,-344 1050.5,-344 1050.5,-302 1193.5,-302 1193.5,-344\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1051,-302 1051,-344 1101,-344 1101,-302 1051,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1056\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">softmax</text>\n",
+       "<text text-anchor=\"start\" x=\"1057.5\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1101,-323 1101,-344 1149,-344 1149,-323 1101,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1111\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1149,-323 1149,-344 1194,-344 1194,-323 1149,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1154\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1101,-302 1101,-323 1149,-323 1149,-302 1101,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1106\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1149,-302 1149,-323 1194,-323 1194,-302 1149,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1154\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 121&#45;&gt;126 -->\n",
+       "<g id=\"edge130\" class=\"edge\">\n",
+       "<title>121&#45;&gt;126</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1122,-379.63C1122,-371.82 1122,-362.73 1122,-354.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1125.5,-354.16 1122,-344.16 1118.5,-354.16 1125.5,-354.16\"/>\n",
+       "</g>\n",
+       "<!-- 140 -->\n",
+       "<g id=\"node141\" class=\"node\">\n",
+       "<title>140</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1436.5,-344 1211.5,-344 1211.5,-302 1436.5,-302 1436.5,-344\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1212,-302 1212,-344 1312,-344 1312,-302 1212,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1217\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">broadcast_tensors</text>\n",
+       "<text text-anchor=\"start\" x=\"1243.5\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1312,-323 1312,-344 1360,-344 1360,-323 1312,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1322\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1360,-323 1360,-344 1437,-344 1437,-323 1360,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1365\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1), (1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1312,-302 1312,-323 1360,-323 1360,-302 1312,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1317\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1360,-302 1360,-323 1437,-323 1437,-302 1360,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1372\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">2 x (1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 121&#45;&gt;140 -->\n",
+       "<g id=\"edge131\" class=\"edge\">\n",
+       "<title>121&#45;&gt;140</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1175.34,-379.93C1201.69,-370.02 1233.65,-357.99 1261.18,-347.64\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1262.59,-350.85 1270.71,-344.05 1260.12,-344.29 1262.59,-350.85\"/>\n",
+       "</g>\n",
+       "<!-- 123 -->\n",
+       "<g id=\"node124\" class=\"node\">\n",
+       "<title>123</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1024.5,-266 881.5,-266 881.5,-224 1024.5,-224 1024.5,-266\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"882,-224 882,-266 932,-266 932,-224 882,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"887\" y=\"-248\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"888.5\" y=\"-237\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"932,-245 932,-266 980,-266 980,-245 932,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"942\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"980,-245 980,-266 1025,-266 1025,-245 980,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"985\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"932,-224 932,-245 980,-245 980,-224 932,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"937\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"980,-224 980,-245 1025,-245 1025,-224 980,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"985\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 122&#45;&gt;123 -->\n",
+       "<g id=\"edge132\" class=\"edge\">\n",
+       "<title>122&#45;&gt;123</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M953,-301.63C953,-293.82 953,-284.73 953,-276.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"956.5,-276.16 953,-266.16 949.5,-276.16 956.5,-276.16\"/>\n",
+       "</g>\n",
+       "<!-- 124 -->\n",
+       "<g id=\"node125\" class=\"node\">\n",
+       "<title>124</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1023,-188 883,-188 883,-146 1023,-146 1023,-188\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"883,-146 883,-188 930,-188 930,-146 883,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"900\" y=\"-170\" font-family=\"Linux libertine\" font-size=\"10.00\">all</text>\n",
+       "<text text-anchor=\"start\" x=\"888\" y=\"-159\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"930,-167 930,-188 978,-188 978,-167 930,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"940\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"978,-167 978,-188 1023,-188 1023,-167 978,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"983\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"930,-146 930,-167 978,-167 978,-146 930,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"935\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"978,-146 978,-167 1023,-167 1023,-146 978,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"990.5\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 123&#45;&gt;124 -->\n",
+       "<g id=\"edge133\" class=\"edge\">\n",
+       "<title>123&#45;&gt;124</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M953,-223.63C953,-215.82 953,-206.73 953,-198.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"956.5,-198.16 953,-188.16 949.5,-198.16 956.5,-198.16\"/>\n",
+       "</g>\n",
+       "<!-- 125 -->\n",
+       "<g id=\"node126\" class=\"node\">\n",
+       "<title>125</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1015.5,-110 890.5,-110 890.5,-68 1015.5,-68 1015.5,-110\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"891,-68 891,-110 938,-110 938,-68 891,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"908\" y=\"-92\" font-family=\"Linux libertine\" font-size=\"10.00\">all</text>\n",
+       "<text text-anchor=\"start\" x=\"896\" y=\"-81\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"938,-89 938,-110 986,-110 986,-89 938,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"948\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"986,-89 986,-110 1016,-110 1016,-89 986,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"991\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"938,-68 938,-89 986,-89 986,-68 938,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"943\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"986,-68 986,-89 1016,-89 1016,-68 986,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"995.5\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">() </text>\n",
+       "</g>\n",
+       "<!-- 124&#45;&gt;125 -->\n",
+       "<g id=\"edge134\" class=\"edge\">\n",
+       "<title>124&#45;&gt;125</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M953,-145.63C953,-137.82 953,-128.73 953,-120.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"956.5,-120.16 953,-110.16 949.5,-120.16 956.5,-120.16\"/>\n",
+       "</g>\n",
+       "<!-- 127 -->\n",
+       "<g id=\"node128\" class=\"node\">\n",
+       "<title>127</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1193.5,-266 1050.5,-266 1050.5,-224 1193.5,-224 1193.5,-266\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1051,-224 1051,-266 1101,-266 1101,-224 1051,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1056\" y=\"-248\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1057.5\" y=\"-237\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1101,-245 1101,-266 1149,-266 1149,-245 1101,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1111\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1149,-245 1149,-266 1194,-266 1194,-245 1149,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1154\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1101,-224 1101,-245 1149,-245 1149,-224 1101,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1106\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1149,-224 1149,-245 1194,-245 1194,-224 1149,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1154\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "</g>\n",
+       "<!-- 126&#45;&gt;127 -->\n",
+       "<g id=\"edge135\" class=\"edge\">\n",
+       "<title>126&#45;&gt;127</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1122,-301.63C1122,-293.82 1122,-284.73 1122,-276.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1125.5,-276.16 1122,-266.16 1118.5,-276.16 1125.5,-276.16\"/>\n",
+       "</g>\n",
+       "<!-- 128 -->\n",
+       "<g id=\"node129\" class=\"node\">\n",
+       "<title>128</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1204.5,-188 1041.5,-188 1041.5,-146 1204.5,-146 1204.5,-188\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1042,-146 1042,-188 1112,-188 1112,-146 1042,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1047\" y=\"-170\" font-family=\"Linux libertine\" font-size=\"10.00\">multinomial</text>\n",
+       "<text text-anchor=\"start\" x=\"1058.5\" y=\"-159\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1112,-167 1112,-188 1160,-188 1160,-167 1112,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1122\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1160,-167 1160,-188 1205,-188 1205,-167 1160,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1165\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1112,-146 1112,-167 1160,-167 1160,-146 1112,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1117\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1160,-146 1160,-167 1205,-167 1205,-146 1160,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1168\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 127&#45;&gt;128 -->\n",
+       "<g id=\"edge136\" class=\"edge\">\n",
+       "<title>127&#45;&gt;128</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1122.27,-223.63C1122.37,-215.82 1122.49,-206.73 1122.6,-198.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1126.1,-198.2 1122.73,-188.16 1119.1,-198.11 1126.1,-198.2\"/>\n",
+       "</g>\n",
+       "<!-- 129 -->\n",
+       "<g id=\"node130\" class=\"node\">\n",
+       "<title>129</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1799,-656 1665,-656 1665,-614 1799,-614 1799,-656\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1665,-614 1665,-656 1712,-656 1712,-614 1665,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1670\" y=\"-638\" font-family=\"Linux libertine\" font-size=\"10.00\">__get__</text>\n",
+       "<text text-anchor=\"start\" x=\"1670\" y=\"-627\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1712,-635 1712,-656 1760,-656 1760,-635 1712,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1722\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1760,-635 1760,-656 1799,-656 1799,-635 1760,-635\"/>\n",
+       "<text text-anchor=\"start\" x=\"1765\" y=\"-643\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1712,-614 1712,-635 1760,-635 1760,-614 1712,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1717\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1760,-614 1760,-635 1799,-635 1799,-614 1760,-614\"/>\n",
+       "<text text-anchor=\"start\" x=\"1765\" y=\"-622\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 130 -->\n",
+       "<g id=\"node131\" class=\"node\">\n",
+       "<title>130</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1800.5,-578 1663.5,-578 1663.5,-536 1800.5,-536 1800.5,-578\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1664,-536 1664,-578 1714,-578 1714,-536 1664,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1669\" y=\"-560\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1670.5\" y=\"-549\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1714,-557 1714,-578 1762,-578 1762,-557 1714,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1724\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1762,-557 1762,-578 1801,-578 1801,-557 1762,-557\"/>\n",
+       "<text text-anchor=\"start\" x=\"1767\" y=\"-565\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1714,-536 1714,-557 1762,-557 1762,-536 1714,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1719\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1762,-536 1762,-557 1801,-557 1801,-536 1762,-536\"/>\n",
+       "<text text-anchor=\"start\" x=\"1771.5\" y=\"-544\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 129&#45;&gt;130 -->\n",
+       "<g id=\"edge137\" class=\"edge\">\n",
+       "<title>129&#45;&gt;130</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1732,-613.63C1732,-605.82 1732,-596.73 1732,-588.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1735.5,-588.16 1732,-578.16 1728.5,-588.16 1735.5,-588.16\"/>\n",
+       "</g>\n",
+       "<!-- 131 -->\n",
+       "<g id=\"node132\" class=\"node\">\n",
+       "<title>131</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1726.5,-500 1587.5,-500 1587.5,-458 1726.5,-458 1726.5,-500\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1588,-458 1588,-500 1649,-500 1649,-458 1588,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1593\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">remainder</text>\n",
+       "<text text-anchor=\"start\" x=\"1600\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1649,-479 1649,-500 1697,-500 1697,-479 1649,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1659\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1697,-479 1697,-500 1727,-500 1727,-479 1697,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1702\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1649,-458 1649,-479 1697,-479 1697,-458 1649,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1654\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1697,-458 1697,-479 1727,-479 1727,-458 1697,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1702\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 130&#45;&gt;131 -->\n",
+       "<g id=\"edge138\" class=\"edge\">\n",
+       "<title>130&#45;&gt;131</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1711.9,-535.63C1703.34,-526.96 1693.23,-516.71 1684.03,-507.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1686.41,-504.82 1676.89,-500.16 1681.43,-509.74 1686.41,-504.82\"/>\n",
+       "</g>\n",
+       "<!-- 133 -->\n",
+       "<g id=\"node134\" class=\"node\">\n",
+       "<title>133</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1869.5,-500 1744.5,-500 1744.5,-458 1869.5,-458 1869.5,-500\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1745,-458 1745,-500 1792,-500 1792,-458 1745,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1762\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">ge</text>\n",
+       "<text text-anchor=\"start\" x=\"1750\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1792,-479 1792,-500 1840,-500 1840,-479 1792,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1802\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1840,-479 1840,-500 1870,-500 1870,-479 1840,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1845\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1792,-458 1792,-479 1840,-479 1840,-458 1792,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1797\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1840,-458 1840,-479 1870,-479 1870,-458 1840,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1845\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 130&#45;&gt;133 -->\n",
+       "<g id=\"edge139\" class=\"edge\">\n",
+       "<title>130&#45;&gt;133</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1752.1,-535.63C1760.66,-526.96 1770.77,-516.71 1779.97,-507.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1782.57,-509.74 1787.11,-500.16 1777.59,-504.82 1782.57,-509.74\"/>\n",
+       "</g>\n",
+       "<!-- 135 -->\n",
+       "<g id=\"node136\" class=\"node\">\n",
+       "<title>135</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1960.5,-422 1835.5,-422 1835.5,-380 1960.5,-380 1960.5,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1836,-380 1836,-422 1883,-422 1883,-380 1836,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1854.5\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">le</text>\n",
+       "<text text-anchor=\"start\" x=\"1841\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1883,-401 1883,-422 1931,-422 1931,-401 1883,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1893\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1931,-401 1931,-422 1961,-422 1961,-401 1931,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1936\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1883,-380 1883,-401 1931,-401 1931,-380 1883,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1888\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1931,-380 1931,-401 1961,-401 1961,-380 1931,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1936\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 130&#45;&gt;135 -->\n",
+       "<g id=\"edge140\" class=\"edge\">\n",
+       "<title>130&#45;&gt;135</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1800.71,-544.58C1828.48,-536.57 1858.52,-522.99 1878,-500 1893.84,-481.3 1898.3,-453.43 1899.1,-432.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1902.6,-432.1 1899.24,-422.06 1895.6,-432.01 1902.6,-432.1\"/>\n",
+       "</g>\n",
+       "<!-- 138 -->\n",
+       "<g id=\"node139\" class=\"node\">\n",
+       "<title>138</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1569.5,-500 1444.5,-500 1444.5,-458 1569.5,-458 1569.5,-500\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1445,-458 1445,-500 1492,-500 1492,-458 1445,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1457.5\" y=\"-482\" font-family=\"Linux libertine\" font-size=\"10.00\">long</text>\n",
+       "<text text-anchor=\"start\" x=\"1450\" y=\"-471\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1492,-479 1492,-500 1540,-500 1540,-479 1492,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1502\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1540,-479 1540,-500 1570,-500 1570,-479 1540,-479\"/>\n",
+       "<text text-anchor=\"start\" x=\"1545\" y=\"-487\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1492,-458 1492,-479 1540,-479 1540,-458 1492,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1497\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1540,-458 1540,-479 1570,-479 1570,-458 1540,-458\"/>\n",
+       "<text text-anchor=\"start\" x=\"1545\" y=\"-466\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 130&#45;&gt;138 -->\n",
+       "<g id=\"edge141\" class=\"edge\">\n",
+       "<title>130&#45;&gt;138</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1672.59,-535.93C1642.98,-525.93 1607.01,-513.78 1576.17,-503.36\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1576.95,-499.93 1566.35,-500.05 1574.71,-506.56 1576.95,-499.93\"/>\n",
+       "</g>\n",
+       "<!-- 132 -->\n",
+       "<g id=\"node133\" class=\"node\">\n",
+       "<title>132</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1746.5,-422 1621.5,-422 1621.5,-380 1746.5,-380 1746.5,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1622,-380 1622,-422 1669,-422 1669,-380 1622,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1639\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">eq</text>\n",
+       "<text text-anchor=\"start\" x=\"1627\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1669,-401 1669,-422 1717,-422 1717,-401 1669,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1679\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1717,-401 1717,-422 1747,-422 1747,-401 1717,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1722\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1669,-380 1669,-401 1717,-401 1717,-380 1669,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1674\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1717,-380 1717,-401 1747,-401 1747,-380 1717,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1722\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 131&#45;&gt;132 -->\n",
+       "<g id=\"edge142\" class=\"edge\">\n",
+       "<title>131&#45;&gt;132</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1664.24,-457.63C1667.07,-449.65 1670.38,-440.33 1673.48,-431.62\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1676.79,-432.75 1676.84,-422.16 1670.19,-430.41 1676.79,-432.75\"/>\n",
+       "</g>\n",
+       "<!-- 134 -->\n",
+       "<g id=\"node135\" class=\"node\">\n",
+       "<title>134</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1870,-344 1724,-344 1724,-302 1870,-302 1870,-344\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1724,-302 1724,-344 1774,-344 1774,-302 1724,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1729\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">__and__</text>\n",
+       "<text text-anchor=\"start\" x=\"1730.5\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1774,-323 1774,-344 1822,-344 1822,-323 1774,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1784\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1822,-323 1822,-344 1870,-344 1870,-323 1822,-323\"/>\n",
+       "<text text-anchor=\"start\" x=\"1827\" y=\"-331\" font-family=\"Linux libertine\" font-size=\"10.00\">2 x (1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1774,-302 1774,-323 1822,-323 1822,-302 1774,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1779\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1822,-302 1822,-323 1870,-323 1870,-302 1822,-302\"/>\n",
+       "<text text-anchor=\"start\" x=\"1836\" y=\"-310\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 132&#45;&gt;134 -->\n",
+       "<g id=\"edge143\" class=\"edge\">\n",
+       "<title>132&#45;&gt;134</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1713.99,-379.83C1727.65,-370.64 1743.99,-359.65 1758.54,-349.86\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1760.53,-352.75 1766.87,-344.26 1756.62,-346.94 1760.53,-352.75\"/>\n",
+       "</g>\n",
+       "<!-- 133&#45;&gt;134 -->\n",
+       "<g id=\"edge144\" class=\"edge\">\n",
+       "<title>133&#45;&gt;134</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1805.69,-457.87C1803.98,-431.47 1800.96,-384.91 1798.96,-354.26\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1802.45,-353.89 1798.31,-344.14 1795.46,-354.34 1802.45,-353.89\"/>\n",
+       "</g>\n",
+       "<!-- 136 -->\n",
+       "<g id=\"node137\" class=\"node\">\n",
+       "<title>136</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1920,-266 1774,-266 1774,-224 1920,-224 1920,-266\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1774,-224 1774,-266 1824,-266 1824,-224 1774,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1779\" y=\"-248\" font-family=\"Linux libertine\" font-size=\"10.00\">__and__</text>\n",
+       "<text text-anchor=\"start\" x=\"1780.5\" y=\"-237\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1824,-245 1824,-266 1872,-266 1872,-245 1824,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1834\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1872,-245 1872,-266 1920,-266 1920,-245 1872,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1877\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">2 x (1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1824,-224 1824,-245 1872,-245 1872,-224 1824,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1829\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1872,-224 1872,-245 1920,-245 1920,-224 1872,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1886\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 134&#45;&gt;136 -->\n",
+       "<g id=\"edge145\" class=\"edge\">\n",
+       "<title>134&#45;&gt;136</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1810.4,-301.63C1815.82,-293.39 1822.18,-283.72 1828.07,-274.78\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1831.17,-276.44 1833.74,-266.16 1825.32,-272.59 1831.17,-276.44\"/>\n",
+       "</g>\n",
+       "<!-- 135&#45;&gt;136 -->\n",
+       "<g id=\"edge146\" class=\"edge\">\n",
+       "<title>135&#45;&gt;136</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1896.06,-379.96C1893.66,-359.78 1888.65,-328 1879,-302 1875.6,-292.83 1870.79,-283.39 1865.93,-274.94\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1868.84,-272.99 1860.7,-266.21 1862.84,-276.59 1868.84,-272.99\"/>\n",
+       "</g>\n",
+       "<!-- 137 -->\n",
+       "<g id=\"node138\" class=\"node\">\n",
+       "<title>137</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1909.5,-188 1784.5,-188 1784.5,-146 1909.5,-146 1909.5,-188\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1785,-146 1785,-188 1832,-188 1832,-146 1785,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1802\" y=\"-170\" font-family=\"Linux libertine\" font-size=\"10.00\">all</text>\n",
+       "<text text-anchor=\"start\" x=\"1790\" y=\"-159\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1832,-167 1832,-188 1880,-188 1880,-167 1832,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1842\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1880,-167 1880,-188 1910,-188 1910,-167 1880,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1885\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1832,-146 1832,-167 1880,-167 1880,-146 1832,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1837\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1880,-146 1880,-167 1910,-167 1910,-146 1880,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1889.5\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">() </text>\n",
+       "</g>\n",
+       "<!-- 136&#45;&gt;137 -->\n",
+       "<g id=\"edge147\" class=\"edge\">\n",
+       "<title>136&#45;&gt;137</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1847,-223.63C1847,-215.82 1847,-206.73 1847,-198.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1850.5,-198.16 1847,-188.16 1843.5,-198.16 1850.5,-198.16\"/>\n",
+       "</g>\n",
+       "<!-- 139 -->\n",
+       "<g id=\"node140\" class=\"node\">\n",
+       "<title>139</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1412,-422 1262,-422 1262,-380 1412,-380 1412,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1262,-380 1262,-422 1325,-422 1325,-380 1262,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1267\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">unsqueeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1275\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1325,-401 1325,-422 1373,-422 1373,-401 1325,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1335\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1373,-401 1373,-422 1412,-422 1412,-401 1373,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1382.5\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1325,-380 1325,-401 1373,-401 1373,-380 1325,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1330\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1373,-380 1373,-401 1412,-401 1412,-380 1373,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1378\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 138&#45;&gt;139 -->\n",
+       "<g id=\"edge148\" class=\"edge\">\n",
+       "<title>138&#45;&gt;139</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1461.89,-457.83C1440.19,-448.13 1414,-436.43 1391.23,-426.24\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1392.58,-423.02 1382.03,-422.13 1389.73,-429.41 1392.58,-423.02\"/>\n",
+       "</g>\n",
+       "<!-- 145 -->\n",
+       "<g id=\"node146\" class=\"node\">\n",
+       "<title>145</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1571,-422 1443,-422 1443,-380 1571,-380 1571,-422\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1443,-380 1443,-422 1493,-422 1493,-380 1443,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1448\" y=\"-404\" font-family=\"Linux libertine\" font-size=\"10.00\">reshape</text>\n",
+       "<text text-anchor=\"start\" x=\"1449.5\" y=\"-393\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-401 1493,-422 1541,-422 1541,-401 1493,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1503\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1541,-401 1541,-422 1571,-422 1571,-401 1541,-401\"/>\n",
+       "<text text-anchor=\"start\" x=\"1546\" y=\"-409\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1493,-380 1493,-401 1541,-401 1541,-380 1493,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1498\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1541,-380 1541,-401 1571,-401 1571,-380 1541,-380\"/>\n",
+       "<text text-anchor=\"start\" x=\"1546\" y=\"-388\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 138&#45;&gt;145 -->\n",
+       "<g id=\"edge149\" class=\"edge\">\n",
+       "<title>138&#45;&gt;145</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1507,-457.63C1507,-449.82 1507,-440.73 1507,-432.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1510.5,-432.16 1507,-422.16 1503.5,-432.16 1510.5,-432.16\"/>\n",
+       "</g>\n",
+       "<!-- 139&#45;&gt;140 -->\n",
+       "<g id=\"edge150\" class=\"edge\">\n",
+       "<title>139&#45;&gt;140</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1333.52,-379.63C1332.18,-371.82 1330.62,-362.73 1329.16,-354.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1332.58,-353.43 1327.45,-344.16 1325.68,-354.61 1332.58,-353.43\"/>\n",
+       "</g>\n",
+       "<!-- 141 -->\n",
+       "<g id=\"node142\" class=\"node\">\n",
+       "<title>141</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1480.5,-266 1317.5,-266 1317.5,-224 1480.5,-224 1480.5,-266\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1318,-224 1318,-266 1388,-266 1388,-224 1318,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1323\" y=\"-248\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
+       "<text text-anchor=\"start\" x=\"1334.5\" y=\"-237\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-245 1388,-266 1436,-266 1436,-245 1388,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1398\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1436,-245 1436,-266 1481,-266 1481,-245 1436,-245\"/>\n",
+       "<text text-anchor=\"start\" x=\"1441\" y=\"-253\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1388,-224 1388,-245 1436,-245 1436,-224 1388,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1393\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1436,-224 1436,-245 1481,-245 1481,-224 1436,-224\"/>\n",
+       "<text text-anchor=\"start\" x=\"1444\" y=\"-232\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 140&#45;&gt;141 -->\n",
+       "<g id=\"edge151\" class=\"edge\">\n",
+       "<title>140&#45;&gt;141</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1344.1,-301.63C1352.66,-292.96 1362.77,-282.71 1371.97,-273.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1374.57,-275.74 1379.11,-266.16 1369.59,-270.82 1374.57,-275.74\"/>\n",
+       "</g>\n",
+       "<!-- 142 -->\n",
+       "<g id=\"node143\" class=\"node\">\n",
+       "<title>142</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1430,-188 1258,-188 1258,-146 1430,-146 1430,-188\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1258,-146 1258,-188 1305,-188 1305,-146 1258,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1265.5\" y=\"-170\" font-family=\"Linux libertine\" font-size=\"10.00\">gather</text>\n",
+       "<text text-anchor=\"start\" x=\"1263\" y=\"-159\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1305,-167 1305,-188 1353,-188 1353,-167 1305,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1315\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-167 1353,-188 1430,-188 1430,-167 1353,-167\"/>\n",
+       "<text text-anchor=\"start\" x=\"1358\" y=\"-175\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 25), (1, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1305,-146 1305,-167 1353,-167 1353,-146 1305,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1310\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1353,-146 1353,-167 1430,-167 1430,-146 1353,-146\"/>\n",
+       "<text text-anchor=\"start\" x=\"1377\" y=\"-154\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "</g>\n",
+       "<!-- 140&#45;&gt;142 -->\n",
+       "<g id=\"edge152\" class=\"edge\">\n",
+       "<title>140&#45;&gt;142</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1315.52,-301.78C1308.32,-281.73 1300.28,-250.3 1308,-224 1310.83,-214.37 1315.88,-204.83 1321.32,-196.43\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1324.25,-198.34 1327.06,-188.13 1318.49,-194.36 1324.25,-198.34\"/>\n",
+       "</g>\n",
+       "<!-- 141&#45;&gt;142 -->\n",
+       "<g id=\"edge153\" class=\"edge\">\n",
+       "<title>141&#45;&gt;142</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1384.26,-223.63C1378.23,-215.3 1371.15,-205.52 1364.63,-196.5\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1367.29,-194.21 1358.59,-188.16 1361.62,-198.31 1367.29,-194.21\"/>\n",
+       "</g>\n",
+       "<!-- 143 -->\n",
+       "<g id=\"node144\" class=\"node\">\n",
+       "<title>143</title>\n",
+       "<polygon fill=\"aliceblue\" stroke=\"transparent\" points=\"1413,-110 1275,-110 1275,-68 1413,-68 1413,-110\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1275,-68 1275,-110 1326,-110 1326,-68 1275,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"1280\" y=\"-92\" font-family=\"Linux libertine\" font-size=\"10.00\">squeeze</text>\n",
+       "<text text-anchor=\"start\" x=\"1282\" y=\"-81\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1326,-89 1326,-110 1374,-110 1374,-89 1326,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"1336\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1374,-89 1374,-110 1413,-110 1413,-89 1374,-89\"/>\n",
+       "<text text-anchor=\"start\" x=\"1379\" y=\"-97\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 1) </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1326,-68 1326,-89 1374,-89 1374,-68 1326,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"1331\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1374,-68 1374,-89 1413,-89 1413,-68 1374,-68\"/>\n",
+       "<text text-anchor=\"start\" x=\"1383.5\" y=\"-76\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,) </text>\n",
+       "</g>\n",
+       "<!-- 142&#45;&gt;143 -->\n",
+       "<g id=\"edge154\" class=\"edge\">\n",
+       "<title>142&#45;&gt;143</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1344,-145.63C1344,-137.82 1344,-128.73 1344,-120.18\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1347.5,-120.16 1344,-110.16 1340.5,-120.16 1347.5,-120.16\"/>\n",
+       "</g>\n",
+       "<!-- 144 -->\n",
+       "<g id=\"node145\" class=\"node\">\n",
+       "<title>144</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"1396,-32 1292,-32 1292,0 1396,0 1396,-32\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1292,0 1292,-32 1369,-32 1369,0 1292,0\"/>\n",
+       "<text text-anchor=\"start\" x=\"1297\" y=\"-19\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"1312\" y=\"-8\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1369,0 1369,-32 1396,-32 1396,0 1369,0\"/>\n",
+       "<text text-anchor=\"start\" x=\"1374\" y=\"-13.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,)</text>\n",
+       "</g>\n",
+       "<!-- 143&#45;&gt;144 -->\n",
+       "<g id=\"edge155\" class=\"edge\">\n",
+       "<title>143&#45;&gt;144</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1344,-67.84C1344,-59.89 1344,-50.66 1344,-42.26\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1347.5,-42.24 1344,-32.24 1340.5,-42.24 1347.5,-42.24\"/>\n",
+       "</g>\n",
+       "<!-- 146 -->\n",
+       "<g id=\"node147\" class=\"node\">\n",
+       "<title>146</title>\n",
+       "<polygon fill=\"lightyellow\" stroke=\"transparent\" points=\"1559,-339 1455,-339 1455,-307 1559,-307 1559,-339\"/>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1455,-307 1455,-339 1532,-339 1532,-307 1455,-307\"/>\n",
+       "<text text-anchor=\"start\" x=\"1460\" y=\"-326\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
+       "<text text-anchor=\"start\" x=\"1475\" y=\"-315\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
+       "<polygon fill=\"none\" stroke=\"black\" points=\"1532,-307 1532,-339 1559,-339 1559,-307 1532,-307\"/>\n",
+       "<text text-anchor=\"start\" x=\"1537\" y=\"-320.5\" font-family=\"Linux libertine\" font-size=\"10.00\">(1,)</text>\n",
+       "</g>\n",
+       "<!-- 145&#45;&gt;146 -->\n",
+       "<g id=\"edge156\" class=\"edge\">\n",
+       "<title>145&#45;&gt;146</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M1507,-379.63C1507,-370.33 1507,-359.22 1507,-349.39\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"1510.5,-349.36 1507,-339.36 1503.5,-349.36 1510.5,-349.36\"/>\n",
+       "</g>\n",
+       "</g>\n",
+       "</svg>\n"
+      ],
+      "text/plain": [
+       "<graphviz.graphs.Digraph at 0x7fc785b5c510>"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_graph.visual_graph"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/temp.ipynb b/examples/temp.ipynb
new file mode 100644
index 00000000..44a1cf02
--- /dev/null
+++ b/examples/temp.ipynb
@@ -0,0 +1,384 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([15, 10])"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from gymnasium.spaces import MultiDiscrete\n",
+    "import numpy as np\n",
+    "action_space = MultiDiscrete(np.array([20,25]), seed=42)\n",
+    "action_space.sample()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([17, 17])"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "action_space.sample()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from nocturne.envs.nocturne_gymnasium import NocturneGymnasium\n",
+    "import yaml\n",
+    "from nocturne.envs.base_env import BaseEnv\n",
+    "from nocturne.envs.vec_env_ma import MultiAgentAsVecEnv\n",
+    "\n",
+    "# Load environment settings\n",
+    "with open(f\"../configs/env_config.yaml\", \"r\") as stream:\n",
+    "    env_config = yaml.safe_load(stream)\n",
+    "\n",
+    "# Initialize environment\n",
+    "env = BaseEnv(config=env_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gymnasiumEnv = NocturneGymnasium(env)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MultiDiscrete([20 25])"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "gymnasiumEnv.action_space"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{6: array([0.30362597, 0.54050583, 0.16309013, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 23: array([0.33037499, 0.583     , 0.15996636, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 2: array([0.29274434, 0.50419724, 0.15200901, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 8: array([0.30145869, 0.53301573, 0.16035738, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 9: array([0.27850392, 0.50491506, 0.16618462, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 1: array([0.28742164, 0.52419186, 0.17056067, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 14: array([0.27822894, 0.51315653, 0.1849147 , ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 18: array([0.28400436, 0.51011646, 0.13745898, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 4: array([0.31240156, 0.52112043, 0.16890202, ..., 0.        , 0.        ,\n",
+       "        0.        ]),\n",
+       " 0: array([0.27941427, 0.51273805, 0.16114239, ..., 0.        , 0.        ,\n",
+       "        0.        ])}"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "gymnasiumEnv.reset()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Reset\n",
+    "obs_dict = gymnasiumEnv.reset()\n",
+    "\n",
+    "# Get info\n",
+    "agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
+    "dead_agent_ids = []\n",
+    "num_agents = len(agent_ids)\n",
+    "rewards = {agent_id: 0 for agent_id in agent_ids}\n",
+    "\n",
+    "for step in range(1000):\n",
+    "\n",
+    "    # Sample actions\n",
+    "    action_dict = {\n",
+    "        agent_id: env.action_space.sample() \n",
+    "        for agent_id in agent_ids\n",
+    "        if agent_id not in dead_agent_ids\n",
+    "    }\n",
+    "    # Step in env\n",
+    "    obs_dict, rew_dict, done_dict, info_dict = gymnasiumEnv.step(action_dict)\n",
+    "\n",
+    "    for agent_id in action_dict.keys():\n",
+    "        rewards[agent_id] += rew_dict[agent_id]\n",
+    "\n",
+    "    # Update dead agents\n",
+    "    for agent_id, is_done in done_dict.items():\n",
+    "        if is_done and agent_id not in dead_agent_ids:\n",
+    "            dead_agent_ids.append(agent_id)\n",
+    "\n",
+    "    # Reset if all agents are done\n",
+    "    if done_dict[\"__all__\"]:\n",
+    "        print(f'Done after {env.step_num} steps -- total return in episode: {rewards}')\n",
+    "        obs_dict = gymnasiumEnv.reset()\n",
+    "        agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
+    "        dead_agent_ids = []\n",
+    "        rewards = {agent_id: 0 for agent_id in agent_ids}\n",
+    "\n",
+    "# Close environment\n",
+    "env.close()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from stable_baselines3.common.vec_env import SubprocVecEnv"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_env(env_config):\n",
+    "    return NocturneGymnasium(BaseEnv(config=env_config)) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "envs = SubprocVecEnv([lambda: make_env(env_config) for _ in range(4)])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Reset\n",
+    "obs_dicts = envs.reset()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent_ids_batch = []\n",
+    "dead_agent_ids_batch = []\n",
+    "num_agents_batch = []\n",
+    "rewards_batch = []\n",
+    "for obs_dict in obs_dicts:\n",
+    "    agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
+    "    dead_agent_ids = []\n",
+    "    num_agents = len(agent_ids)\n",
+    "    rewards = {agent_id: 0 for agent_id in agent_ids}\n",
+    "    agent_ids_batch.append(agent_ids)\n",
+    "    dead_agent_ids_batch.append(dead_agent_ids)\n",
+    "    num_agents_batch.append(num_agents)\n",
+    "    rewards_batch.append(rewards)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{6: 6, 23: 18, 2: 24, 8: 1, 9: 2, 1: 18, 14: 1, 18: 13, 4: 11, 0: 18},\n",
+       " {6: 0, 23: 14, 2: 8, 8: 22, 9: 1, 1: 4, 14: 18, 18: 18, 4: 1, 0: 23},\n",
+       " {6: 9, 23: 22, 2: 1, 8: 2, 9: 0, 1: 20, 14: 5, 18: 14, 4: 7, 0: 13},\n",
+       " {6: 13, 23: 1, 2: 8, 8: 13, 9: 9, 1: 16, 14: 9, 18: 9, 4: 12, 0: 23}]"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "action_dicts = [\n",
+    "        {\n",
+    "            agent_id: env.action_space.sample() \n",
+    "            for agent_id in agent_ids\n",
+    "            if agent_id not in dead_agent_ids\n",
+    "        }\n",
+    "        for agent_ids, dead_agent_ids in zip(agent_ids_batch, dead_agent_ids_batch)\n",
+    "    ]\n",
+    "action_dicts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "KeyError",
+     "evalue": "7",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[9], line 17\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m rew_dict, rewards \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(rew_dicts, rewards_batch):\n\u001b[1;32m     16\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m agent_id \u001b[38;5;129;01min\u001b[39;00m rew_dict\u001b[38;5;241m.\u001b[39mkeys():\n\u001b[0;32m---> 17\u001b[0m         \u001b[43mrewards\u001b[49m\u001b[43m[\u001b[49m\u001b[43magent_id\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m rew_dict[agent_id] \n\u001b[1;32m     19\u001b[0m \u001b[38;5;66;03m# Update dead agents\u001b[39;00m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m done_dict, dead_agent_ids \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(done_dicts, dead_agent_ids_batch):\n",
+      "\u001b[0;31mKeyError\u001b[0m: 7"
+     ]
+    }
+   ],
+   "source": [
+    "for step in range(1000):\n",
+    "\n",
+    "    # Sample actions\n",
+    "    action_dicts = [\n",
+    "        {\n",
+    "            agent_id: env.action_space.sample() \n",
+    "            for agent_id in agent_ids\n",
+    "            if agent_id not in dead_agent_ids\n",
+    "        }\n",
+    "        for agent_ids, dead_agent_ids in zip(agent_ids_batch, dead_agent_ids_batch)\n",
+    "    ]\n",
+    "    # Step in env\n",
+    "    obs_dicts, rew_dicts, done_dicts, info_dicts = envs.step(action_dicts)\n",
+    "\n",
+    "    for rew_dict, rewards in zip(rew_dicts, rewards_batch):\n",
+    "        for agent_id in rew_dict.keys():\n",
+    "            rewards[agent_id] += rew_dict[agent_id] \n",
+    "    \n",
+    "    # Update dead agents\n",
+    "    for done_dict, dead_agent_ids in zip(done_dicts, dead_agent_ids_batch):\n",
+    "        for agent_id, is_done in done_dict.items():\n",
+    "            if is_done and agent_id not in dead_agent_ids:\n",
+    "                dead_agent_ids.append(agent_id)\n",
+    "\n",
+    "    # Reset if all agents are done\n",
+    "    if all([done_dict[\"__all__\"] for done_dict in done_dicts]):\n",
+    "        print(f'Done after {env.step_num} steps -- total return in episode: {rewards}')\n",
+    "        obs_dicts = envs.reset()\n",
+    "        agent_ids_batch = []\n",
+    "        dead_agent_ids_batch = []\n",
+    "        num_agents_batch = []\n",
+    "        rewards_batch = []\n",
+    "        for obs_dict in obs_dicts:\n",
+    "            agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
+    "            dead_agent_ids = []\n",
+    "            num_agents = len(agent_ids)\n",
+    "            rewards = {agent_id: 0 for agent_id in agent_ids}\n",
+    "            agent_ids_batch.append(agent_ids)\n",
+    "            dead_agent_ids_batch.append(dead_agent_ids)\n",
+    "            num_agents_batch.append(num_agents)\n",
+    "            rewards_batch.append(rewards)\n",
+    "\n",
+    "    # # Sample actions\n",
+    "    # action_dict = {\n",
+    "    #     agent_id: env.action_space.sample() \n",
+    "    #     for agent_id in agent_ids\n",
+    "    #     if agent_id not in dead_agent_ids\n",
+    "    # }\n",
+    "    # # Step in env\n",
+    "    # obs_dict, rew_dict, done_dict, info_dict = gymnasiumEnv.step(action_dict)\n",
+    "\n",
+    "    # for agent_id in action_dict.keys():\n",
+    "    #     rewards[agent_id] += rew_dict[agent_id]\n",
+    "\n",
+    "    # # Update dead agents\n",
+    "    # for agent_id, is_done in done_dict.items():\n",
+    "    #     if is_done and agent_id not in dead_agent_ids:\n",
+    "    #         dead_agent_ids.append(agent_id)\n",
+    "\n",
+    "    # # Reset if all agents are done\n",
+    "    # if done_dict[\"__all__\"]:\n",
+    "    #     print(f'Done after {env.step_num} steps -- total return in episode: {rewards}')\n",
+    "    #     obs_dict = gymnasiumEnv.reset()\n",
+    "    #     agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
+    "    #     dead_agent_ids = []\n",
+    "    #     rewards = {agent_id: 0 for agent_id in agent_ids}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/nocturne/envs/nocturne_gymnasium.py b/nocturne/envs/nocturne_gymnasium.py
new file mode 100644
index 00000000..099031b2
--- /dev/null
+++ b/nocturne/envs/nocturne_gymnasium.py
@@ -0,0 +1,136 @@
+"""Gymnasium vectorizable environment wrapper for Nocturne."""
+import logging
+import time
+from copy import deepcopy
+from typing import Any, Dict, List, TypeVar
+
+import gym
+import gymnasium
+import numpy as np
+
+from nocturne.envs.base_env import BaseEnv
+from utils.config import load_config
+
+logging.basicConfig(level=logging.INFO)
+
+ObsType = TypeVar("ObsType")
+ActType = TypeVar("ActType")
+
+class NocturneGymnasium(gymnasium.Env):
+    """Nocturne environment wrapper for compatible with SB3.
+    """
+
+    def __init__(self, config, num_agents, psr=False):
+        self.env = BaseEnv(config)
+
+        # Make action and observation spaces compatible with SB3 (requires gymnasium)
+        self.action_space = gymnasium.spaces.MultiDiscrete([self.env.config.max_num_vehicles, self.env.action_space.n])
+        self.observation_space = gym.spaces.Box(-np.inf, np.inf, self.env.observation_space.shape, np.float32)
+        self.num_agents = num_agents  # The maximum number of agents allowed in the environmen
+        self.psr = psr # Whether to use PSR or not
+
+        self.psr_dict = self.init_scene_dict() if psr else None # Initialize dict to keep track of the average reward obtained in each scene
+        self.n_episodes = 0
+        self.episode_lengths = []
+        self.rewards = []  # Log reward per step
+        self.dead_agent_ids = []  # Log dead agents per step
+        self.num_agents_collided = 0  # Keep track of how many agents collided
+        self.total_agents_in_rollout = 0 # Log total number of agents in rollout
+        self.num_agents_goal_achieved = 0 # Keep track of how many agents reached their goal
+        self.agents_in_scene = []
+        self.filename = None # If provided, always use the same file 
+
+    def step(self, actions):
+        """Take a step in the environment, convert dicts to np arrays.
+
+        Args
+        ----
+            action (Dict): Dictionary with a single action for the controlled vehicle.
+
+        Returns
+        -------
+            observation, reward, terminated, truncated, info (np.ndarray, float, bool, bool, dict)
+        """
+        next_obs_dict, rewards_dict, dones_dict, info_dict = self.env.step(
+            action_dict=actions
+        )
+
+        return (
+            next_obs_dict,
+            rewards_dict,
+            dones_dict,
+            False,
+            info_dict,
+        )
+
+    def reset(self, seed=None):
+        """Reset environment and return initial observations."""
+        obs_dict = self.env.reset()
+                
+        # Reset Nocturne env
+        obs_dict = self.env.reset(self.filename, self.psr_dict)
+
+        # Reset storage
+        self.agent_ids = []
+        self.rewards = []
+        self.dead_agent_ids = []
+        self.ep_collisions = 0
+        self.ep_goal_achived = 0
+
+        obs_all = np.full(fill_value=np.nan, shape=(self.num_envs, self.env.observation_space.shape[0]))
+        for idx, agent_id in enumerate(obs_dict.keys()):
+            self.agent_ids.append(agent_id)
+            obs_all[idx, :] = obs_dict[agent_id]
+
+        # Save obs in buffer
+        self._save_obs(obs_all)
+
+        logging.debug(f"RESET - agent ids: {self.agent_ids}")
+
+        # Make dict for storing the last info set for each agent
+        self.last_info_dicts = {agent_id: {} for agent_id in self.agent_ids}
+
+        return self._obs_from_buf(), {}
+
+    def _obs_from_buf(self) -> ObsType:
+        """Get observation from buffer."""
+        return np.copy(self.buf_obs)
+
+    @property
+    def action_space(self):
+        return self.env.action_space
+
+    @action_space.setter
+    def action_space(self, action_space):
+        self.env.action_space = action_space
+
+    @property
+    def observation_space(self):
+        return self.env.observation_space
+
+    @observation_space.setter
+    def observation_space(self, observation_space):
+        self.env.observation_space = observation_space
+
+    def render(self):
+        pass
+
+    def close(self):
+        pass
+
+    @property
+    def seed(self, seed=None):
+        return None
+
+    @seed.setter
+    def seed(self, seed=None):
+        pass
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+    def get_attr(self, attr_name: str):
+        return getattr(self._env, attr_name)
+
+    def set_attr(self, attr_name: str):
+        setattr(self._env, attr_name)

From 33b65ea795e779215194c7cf7666e5a2489b117a Mon Sep 17 00:00:00 2001
From: Aarav Pandya <ap7641@nyu.edu>
Date: Sun, 14 Jan 2024 19:54:00 -0500
Subject: [PATCH 2/2] Pufferlib vectorization support

---
 examples/07_nocturne_pufferlib.ipynb | 169 ++++++++++++
 examples/temp.ipynb                  | 384 ---------------------------
 nocturne/envs/nocturne_gymnasium.py  | 225 +++++++++++++---
 3 files changed, 352 insertions(+), 426 deletions(-)
 create mode 100644 examples/07_nocturne_pufferlib.ipynb
 delete mode 100644 examples/temp.ipynb

diff --git a/examples/07_nocturne_pufferlib.ipynb b/examples/07_nocturne_pufferlib.ipynb
new file mode 100644
index 00000000..ac23d938
--- /dev/null
+++ b/examples/07_nocturne_pufferlib.ipynb
@@ -0,0 +1,169 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from nocturne.envs.nocturne_gymnasium import NocturneGymnasium, CustomPostprocessor\n",
+    "import yaml\n",
+    "from nocturne.envs.base_env import BaseEnv\n",
+    "from nocturne.envs.vec_env_ma import MultiAgentAsVecEnv\n",
+    "\n",
+    "import pufferlib.vectorization\n",
+    "vec = pufferlib.vectorization.Multiprocessing\n",
+    "# vec = pufferlib.vectorization.Serial\n",
+    "\n",
+    "import pufferlib.emulation\n",
+    "import pufferlib.wrappers\n",
+    "\n",
+    "from time import perf_counter\n",
+    "\n",
+    "# Load environment settings\n",
+    "with open(f\"../configs/env_config.yaml\", \"r\") as stream:\n",
+    "    env_config = yaml.safe_load(stream)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_env(env_config):\n",
+    "    return NocturneGymnasium(config=env_config, num_agents=env_config[\"max_num_vehicles\"]) \n",
+    "\n",
+    "def nocturne_creator(env_config):\n",
+    "    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=make_env, env_args=(env_config,), postprocessor_cls=CustomPostprocessor)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# envs = vec(nocturne_creator,env_args=[env_config,], num_envs=4, envs_per_worker=2, env_pool=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# envs.async_reset()\n",
+    "# obs = envs.recv()[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# actions = [envs.single_action_space.sample() for _ in range(4)]\n",
+    "# envs.step(actions)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Average time for single env step: 0.002120851782616228\n",
+      "Average FPS for single env step: 471.50866844944056\n"
+     ]
+    }
+   ],
+   "source": [
+    "NUM_STEPS = 1000\n",
+    "\n",
+    "env = make_env(env_config)\n",
+    "env.reset()\n",
+    "\n",
+    "total_time = 0\n",
+    "\n",
+    "for i in range(NUM_STEPS):\n",
+    "    actions = env.action_space.sample()\n",
+    "    start = perf_counter()\n",
+    "    env.step(actions)\n",
+    "    end = perf_counter()\n",
+    "    total_time += end - start\n",
+    "\n",
+    "print(f\"Average time for single env step: {total_time/NUM_STEPS}\")\n",
+    "print(f\"Average FPS for single env step: {NUM_STEPS/(total_time)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Average time for 1000 env step: 0.0010244677745140506\n",
+      "Average FPS for 1000 env step: 976.1165991525143\n"
+     ]
+    }
+   ],
+   "source": [
+    "NUM_STEPS = 1000\n",
+    "NUM_PROCESSES = 32\n",
+    "NUM_ENVS = 128\n",
+    "assert(NUM_ENVS % NUM_PROCESSES == 0)\n",
+    "NUM_ENVS_PER_WORKER = NUM_ENVS // NUM_PROCESSES\n",
+    "\n",
+    "envs = vec(nocturne_creator,env_args=[env_config,], num_envs=NUM_ENVS, envs_per_worker=1, env_pool=False)\n",
+    "envs.async_reset()\n",
+    "obs = envs.recv()[0]\n",
+    "\n",
+    "total_time = 0\n",
+    "for i in range(NUM_STEPS):\n",
+    "    actions = [envs.single_action_space.sample() for _ in range(NUM_ENVS)]\n",
+    "    start = perf_counter()\n",
+    "    envs.step(actions)\n",
+    "    end = perf_counter()\n",
+    "    total_time += end - start\n",
+    "envs.close()\n",
+    "print(f\"Average time for {NUM_STEPS} env step: {total_time/(NUM_STEPS*NUM_ENVS)}\")\n",
+    "print(f\"Average FPS for {NUM_STEPS} env step: {(NUM_STEPS/(total_time))*NUM_ENVS}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/temp.ipynb b/examples/temp.ipynb
deleted file mode 100644
index 44a1cf02..00000000
--- a/examples/temp.ipynb
+++ /dev/null
@@ -1,384 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "array([15, 10])"
-      ]
-     },
-     "execution_count": 1,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from gymnasium.spaces import MultiDiscrete\n",
-    "import numpy as np\n",
-    "action_space = MultiDiscrete(np.array([20,25]), seed=42)\n",
-    "action_space.sample()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "array([17, 17])"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "action_space.sample()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from nocturne.envs.nocturne_gymnasium import NocturneGymnasium\n",
-    "import yaml\n",
-    "from nocturne.envs.base_env import BaseEnv\n",
-    "from nocturne.envs.vec_env_ma import MultiAgentAsVecEnv\n",
-    "\n",
-    "# Load environment settings\n",
-    "with open(f\"../configs/env_config.yaml\", \"r\") as stream:\n",
-    "    env_config = yaml.safe_load(stream)\n",
-    "\n",
-    "# Initialize environment\n",
-    "env = BaseEnv(config=env_config)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "gymnasiumEnv = NocturneGymnasium(env)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "MultiDiscrete([20 25])"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "gymnasiumEnv.action_space"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{6: array([0.30362597, 0.54050583, 0.16309013, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 23: array([0.33037499, 0.583     , 0.15996636, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 2: array([0.29274434, 0.50419724, 0.15200901, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 8: array([0.30145869, 0.53301573, 0.16035738, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 9: array([0.27850392, 0.50491506, 0.16618462, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 1: array([0.28742164, 0.52419186, 0.17056067, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 14: array([0.27822894, 0.51315653, 0.1849147 , ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 18: array([0.28400436, 0.51011646, 0.13745898, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 4: array([0.31240156, 0.52112043, 0.16890202, ..., 0.        , 0.        ,\n",
-       "        0.        ]),\n",
-       " 0: array([0.27941427, 0.51273805, 0.16114239, ..., 0.        , 0.        ,\n",
-       "        0.        ])}"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "gymnasiumEnv.reset()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Reset\n",
-    "obs_dict = gymnasiumEnv.reset()\n",
-    "\n",
-    "# Get info\n",
-    "agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
-    "dead_agent_ids = []\n",
-    "num_agents = len(agent_ids)\n",
-    "rewards = {agent_id: 0 for agent_id in agent_ids}\n",
-    "\n",
-    "for step in range(1000):\n",
-    "\n",
-    "    # Sample actions\n",
-    "    action_dict = {\n",
-    "        agent_id: env.action_space.sample() \n",
-    "        for agent_id in agent_ids\n",
-    "        if agent_id not in dead_agent_ids\n",
-    "    }\n",
-    "    # Step in env\n",
-    "    obs_dict, rew_dict, done_dict, info_dict = gymnasiumEnv.step(action_dict)\n",
-    "\n",
-    "    for agent_id in action_dict.keys():\n",
-    "        rewards[agent_id] += rew_dict[agent_id]\n",
-    "\n",
-    "    # Update dead agents\n",
-    "    for agent_id, is_done in done_dict.items():\n",
-    "        if is_done and agent_id not in dead_agent_ids:\n",
-    "            dead_agent_ids.append(agent_id)\n",
-    "\n",
-    "    # Reset if all agents are done\n",
-    "    if done_dict[\"__all__\"]:\n",
-    "        print(f'Done after {env.step_num} steps -- total return in episode: {rewards}')\n",
-    "        obs_dict = gymnasiumEnv.reset()\n",
-    "        agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
-    "        dead_agent_ids = []\n",
-    "        rewards = {agent_id: 0 for agent_id in agent_ids}\n",
-    "\n",
-    "# Close environment\n",
-    "env.close()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from stable_baselines3.common.vec_env import SubprocVecEnv"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def make_env(env_config):\n",
-    "    return NocturneGymnasium(BaseEnv(config=env_config)) "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "envs = SubprocVecEnv([lambda: make_env(env_config) for _ in range(4)])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Reset\n",
-    "obs_dicts = envs.reset()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "agent_ids_batch = []\n",
-    "dead_agent_ids_batch = []\n",
-    "num_agents_batch = []\n",
-    "rewards_batch = []\n",
-    "for obs_dict in obs_dicts:\n",
-    "    agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
-    "    dead_agent_ids = []\n",
-    "    num_agents = len(agent_ids)\n",
-    "    rewards = {agent_id: 0 for agent_id in agent_ids}\n",
-    "    agent_ids_batch.append(agent_ids)\n",
-    "    dead_agent_ids_batch.append(dead_agent_ids)\n",
-    "    num_agents_batch.append(num_agents)\n",
-    "    rewards_batch.append(rewards)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[{6: 6, 23: 18, 2: 24, 8: 1, 9: 2, 1: 18, 14: 1, 18: 13, 4: 11, 0: 18},\n",
-       " {6: 0, 23: 14, 2: 8, 8: 22, 9: 1, 1: 4, 14: 18, 18: 18, 4: 1, 0: 23},\n",
-       " {6: 9, 23: 22, 2: 1, 8: 2, 9: 0, 1: 20, 14: 5, 18: 14, 4: 7, 0: 13},\n",
-       " {6: 13, 23: 1, 2: 8, 8: 13, 9: 9, 1: 16, 14: 9, 18: 9, 4: 12, 0: 23}]"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "action_dicts = [\n",
-    "        {\n",
-    "            agent_id: env.action_space.sample() \n",
-    "            for agent_id in agent_ids\n",
-    "            if agent_id not in dead_agent_ids\n",
-    "        }\n",
-    "        for agent_ids, dead_agent_ids in zip(agent_ids_batch, dead_agent_ids_batch)\n",
-    "    ]\n",
-    "action_dicts"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "KeyError",
-     "evalue": "7",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[9], line 17\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m rew_dict, rewards \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(rew_dicts, rewards_batch):\n\u001b[1;32m     16\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m agent_id \u001b[38;5;129;01min\u001b[39;00m rew_dict\u001b[38;5;241m.\u001b[39mkeys():\n\u001b[0;32m---> 17\u001b[0m         \u001b[43mrewards\u001b[49m\u001b[43m[\u001b[49m\u001b[43magent_id\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m rew_dict[agent_id] \n\u001b[1;32m     19\u001b[0m \u001b[38;5;66;03m# Update dead agents\u001b[39;00m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m done_dict, dead_agent_ids \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(done_dicts, dead_agent_ids_batch):\n",
-      "\u001b[0;31mKeyError\u001b[0m: 7"
-     ]
-    }
-   ],
-   "source": [
-    "for step in range(1000):\n",
-    "\n",
-    "    # Sample actions\n",
-    "    action_dicts = [\n",
-    "        {\n",
-    "            agent_id: env.action_space.sample() \n",
-    "            for agent_id in agent_ids\n",
-    "            if agent_id not in dead_agent_ids\n",
-    "        }\n",
-    "        for agent_ids, dead_agent_ids in zip(agent_ids_batch, dead_agent_ids_batch)\n",
-    "    ]\n",
-    "    # Step in env\n",
-    "    obs_dicts, rew_dicts, done_dicts, info_dicts = envs.step(action_dicts)\n",
-    "\n",
-    "    for rew_dict, rewards in zip(rew_dicts, rewards_batch):\n",
-    "        for agent_id in rew_dict.keys():\n",
-    "            rewards[agent_id] += rew_dict[agent_id] \n",
-    "    \n",
-    "    # Update dead agents\n",
-    "    for done_dict, dead_agent_ids in zip(done_dicts, dead_agent_ids_batch):\n",
-    "        for agent_id, is_done in done_dict.items():\n",
-    "            if is_done and agent_id not in dead_agent_ids:\n",
-    "                dead_agent_ids.append(agent_id)\n",
-    "\n",
-    "    # Reset if all agents are done\n",
-    "    if all([done_dict[\"__all__\"] for done_dict in done_dicts]):\n",
-    "        print(f'Done after {env.step_num} steps -- total return in episode: {rewards}')\n",
-    "        obs_dicts = envs.reset()\n",
-    "        agent_ids_batch = []\n",
-    "        dead_agent_ids_batch = []\n",
-    "        num_agents_batch = []\n",
-    "        rewards_batch = []\n",
-    "        for obs_dict in obs_dicts:\n",
-    "            agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
-    "            dead_agent_ids = []\n",
-    "            num_agents = len(agent_ids)\n",
-    "            rewards = {agent_id: 0 for agent_id in agent_ids}\n",
-    "            agent_ids_batch.append(agent_ids)\n",
-    "            dead_agent_ids_batch.append(dead_agent_ids)\n",
-    "            num_agents_batch.append(num_agents)\n",
-    "            rewards_batch.append(rewards)\n",
-    "\n",
-    "    # # Sample actions\n",
-    "    # action_dict = {\n",
-    "    #     agent_id: env.action_space.sample() \n",
-    "    #     for agent_id in agent_ids\n",
-    "    #     if agent_id not in dead_agent_ids\n",
-    "    # }\n",
-    "    # # Step in env\n",
-    "    # obs_dict, rew_dict, done_dict, info_dict = gymnasiumEnv.step(action_dict)\n",
-    "\n",
-    "    # for agent_id in action_dict.keys():\n",
-    "    #     rewards[agent_id] += rew_dict[agent_id]\n",
-    "\n",
-    "    # # Update dead agents\n",
-    "    # for agent_id, is_done in done_dict.items():\n",
-    "    #     if is_done and agent_id not in dead_agent_ids:\n",
-    "    #         dead_agent_ids.append(agent_id)\n",
-    "\n",
-    "    # # Reset if all agents are done\n",
-    "    # if done_dict[\"__all__\"]:\n",
-    "    #     print(f'Done after {env.step_num} steps -- total return in episode: {rewards}')\n",
-    "    #     obs_dict = gymnasiumEnv.reset()\n",
-    "    #     agent_ids = [agent_id for agent_id in obs_dict.keys()]\n",
-    "    #     dead_agent_ids = []\n",
-    "    #     rewards = {agent_id: 0 for agent_id in agent_ids}"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.5"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/nocturne/envs/nocturne_gymnasium.py b/nocturne/envs/nocturne_gymnasium.py
index 099031b2..9dbe76d2 100644
--- a/nocturne/envs/nocturne_gymnasium.py
+++ b/nocturne/envs/nocturne_gymnasium.py
@@ -2,12 +2,14 @@
 import logging
 import time
 from copy import deepcopy
-from typing import Any, Dict, List, TypeVar
+from typing import Any, Dict, List, TypeVar, SupportsFloat
 
-import gym
 import gymnasium
 import numpy as np
 
+from pufferlib.emulation import Postprocessor
+
+
 from nocturne.envs.base_env import BaseEnv
 from utils.config import load_config
 
@@ -24,11 +26,13 @@ def __init__(self, config, num_agents, psr=False):
         self.env = BaseEnv(config)
 
         # Make action and observation spaces compatible with SB3 (requires gymnasium)
-        self.action_space = gymnasium.spaces.MultiDiscrete([self.env.config.max_num_vehicles, self.env.action_space.n])
-        self.observation_space = gym.spaces.Box(-np.inf, np.inf, self.env.observation_space.shape, np.float32)
+        # self.action_space = gymnasium.spaces.MultiDiscrete([self.env.config.max_num_vehicles, self.env.action_space.n])
         self.num_agents = num_agents  # The maximum number of agents allowed in the environmen
+        self.action_space = gymnasium.spaces.MultiDiscrete([self.env.action_space.n] * self.num_agents)
+        self.observation_space = gymnasium.spaces.Box(-np.inf, np.inf, [self.num_agents, self.env.observation_space.shape[0]], np.float32)
         self.psr = psr # Whether to use PSR or not
 
+        self.buf_obs = None  # type: ObsType
         self.psr_dict = self.init_scene_dict() if psr else None # Initialize dict to keep track of the average reward obtained in each scene
         self.n_episodes = 0
         self.episode_lengths = []
@@ -40,33 +44,106 @@ def __init__(self, config, num_agents, psr=False):
         self.agents_in_scene = []
         self.filename = None # If provided, always use the same file 
 
-    def step(self, actions):
-        """Take a step in the environment, convert dicts to np arrays.
+    def step(self, actions) -> tuple[ObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        """Convert action vector to dict and call env.step()."""
 
-        Args
-        ----
-            action (Dict): Dictionary with a single action for the controlled vehicle.
+        agent_actions = {
+            agent_id: actions[idx] for idx, agent_id in enumerate(self.agent_ids) if agent_id not in self.dead_agent_ids
+        }
 
-        Returns
-        -------
-            observation, reward, terminated, truncated, info (np.ndarray, float, bool, bool, dict)
-        """
-        next_obs_dict, rewards_dict, dones_dict, info_dict = self.env.step(
-            action_dict=actions
-        )
+        # Take a step to obtain dicts
+        next_obses_dict, rew_dict, done_dict, info_dict = self.env.step(agent_actions)
+
+        # Update dead agents based on most recent done_dict
+        for agent_id, is_done in done_dict.items():
+            if is_done and agent_id not in self.dead_agent_ids:
+                self.dead_agent_ids.append(agent_id)
+                # Store agents' last info dict
+                self.last_info_dicts[agent_id] = info_dict[agent_id].copy()
+
+        # Storage
+        obs = np.full(fill_value=np.nan, shape=self.observation_space.shape)
+        self.buf_dones = np.full(fill_value=np.nan, shape=(self.num_agents,))
+        self.buf_rews = np.full_like(self.buf_dones, fill_value=np.nan)
+        self.buf_infos = [{} for _ in range(self.num_agents)]
+        
+        # Override NaN placeholder for each agent that is alive
+        for idx, key in enumerate(self.agent_ids):
+            if key in next_obses_dict:
+                self.buf_rews[idx] = rew_dict[key]
+                self.buf_dones[idx] = done_dict[key] * 1
+                self.buf_infos[idx] = info_dict[key]
+                obs[idx, :] = next_obses_dict[key] 
+
+        # Save step reward obtained across all agents
+        self.rewards.append(sum(rew_dict.values()))
+        self.agents_in_scene.append(len(self.agent_ids))
+
+        # Store observation
+        self._save_obs(obs)
+
+        # Reset episode if ALL agents are done
+        if done_dict["__all__"]:
+            for agent_id in self.agent_ids:
+                self.ep_collisions += self.last_info_dicts[agent_id]["collided"] * 1
+                self.ep_goal_achived += self.last_info_dicts[agent_id]["goal_achieved"] * 1
+
+            # Store the fraction of agents that collided in episode
+            self.num_agents_collided += self.ep_collisions
+            self.num_agents_goal_achieved += self.ep_goal_achived
+            self.total_agents_in_rollout += len(self.agent_ids)
+
+            # Save final observation where user can get it, then reset
+            for idx in range(len(self.agent_ids)):
+                self.buf_infos[idx]["terminal_observation"] = obs[idx]
+
+            # Log episode stats
+            ep_len = self.step_num
+            self.n_episodes += 1
+            self.episode_lengths.append(ep_len)
+
+            # Store reward at scene level
+            if self.psr:
+                self.psr_dict[self.env.file]["count"] += 1
+                self.psr_dict[self.env.file]["reward"] += (sum(rew_dict.values())) / len(self.agent_ids)
+                self.psr_dict[self.env.file]["goal_rate"] += self.ep_goal_achived / len(self.agent_ids)
+
+            # Reset
+            obs = self.reset()
 
         return (
-            next_obs_dict,
-            rewards_dict,
-            dones_dict,
+            self._obs_from_buf(),
+            np.copy(self.buf_rews),
+            self.buf_dones.all(),
             False,
-            info_dict,
+            {'infos': deepcopy(self.buf_infos)},
         )
 
+    # def step(self, actions) -> tuple[ObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+    #     """Take a step in the environment, convert dicts to np arrays.
+
+    #     Args
+    #     ----
+    #         action (Dict): Dictionary with a single action for the controlled vehicle.
+
+    #     Returns
+    #     -------
+    #         observation, reward, terminated, truncated, info (np.ndarray, float, bool, bool, dict)
+    #     """
+    #     next_obs_dict, rewards_dict, dones_dict, info_dict = self.env.step(
+    #         action_dict=actions
+    #     )
+
+    #     return (
+    #         next_obs_dict,
+    #         rewards_dict,
+    #         dones_dict,
+    #         False,
+    #         info_dict,
+    #     )
+
     def reset(self, seed=None):
         """Reset environment and return initial observations."""
-        obs_dict = self.env.reset()
-                
         # Reset Nocturne env
         obs_dict = self.env.reset(self.filename, self.psr_dict)
 
@@ -77,7 +154,7 @@ def reset(self, seed=None):
         self.ep_collisions = 0
         self.ep_goal_achived = 0
 
-        obs_all = np.full(fill_value=np.nan, shape=(self.num_envs, self.env.observation_space.shape[0]))
+        obs_all = np.full(fill_value=-np.pi*1e7, shape=self.observation_space.shape, dtype=self.observation_space.dtype)
         for idx, agent_id in enumerate(obs_dict.keys()):
             self.agent_ids.append(agent_id)
             obs_all[idx, :] = obs_dict[agent_id]
@@ -96,21 +173,14 @@ def _obs_from_buf(self) -> ObsType:
         """Get observation from buffer."""
         return np.copy(self.buf_obs)
 
-    @property
-    def action_space(self):
-        return self.env.action_space
-
-    @action_space.setter
-    def action_space(self, action_space):
-        self.env.action_space = action_space
+    def _save_obs(self, obs: ObsType) -> None:
+        """Save observations into buffer."""
+        self.buf_obs = obs
 
     @property
-    def observation_space(self):
-        return self.env.observation_space
-
-    @observation_space.setter
-    def observation_space(self, observation_space):
-        self.env.observation_space = observation_space
+    def step_num(self) -> List[int]:
+        """The episodic timestep."""
+        return self.env.step_num
 
     def render(self):
         pass
@@ -126,11 +196,82 @@ def seed(self, seed=None):
     def seed(self, seed=None):
         pass
 
-    def __getattr__(self, name):
-        return getattr(self._env, name)
+    def get_attr(self, attr_name, indices=None):
+        raise NotImplementedError()
+
+    def set_attr(self, attr_name, value, indices=None) -> None:
+        raise NotImplementedError()
+
+class CustomPostprocessor(Postprocessor):
+    '''Basic postprocessor that injects returns and lengths information into infos and
+    provides an option to pad to a maximum episode length. Works for single-agent and
+    team-based multi-agent environments'''
+    def reset(self, obs):
+        self.epoch_return = 0
+        self.epoch_length = 0
+        self.done = False
+
+    def reward_done_truncated_info(self, reward, done, truncated, info):
+        if isinstance(reward, (list, np.ndarray)):
+            reward = sum(reward)
+
+        # Env is done
+        if self.done:
+            return reward, done, truncated, info
+
+        self.epoch_length += 1
+        self.epoch_return += reward
+
+        if done.all() or truncated:
+            info['return'] = self.epoch_return
+            info['length'] = self.epoch_length
+            self.done = True
+
+        return reward, done, truncated, info
+
+def make_env(env_config, num_agents):
+    return NocturneGymnasium(config=env_config, num_agents=num_agents)
+
+
+def nocturne_creator(env_config, num_agents):
+    return pufferlib.emulation.GymnasiumPufferEnv(env_creator=make_env, env_args=(env_config,num_agents,), postprocessor_cls=CustomPostprocessor)
+
+if __name__ == "__main__":
+    MAX_AGENTS = 3
+    NUM_STEPS = 400
+
+    # Load environment variables and config
+    env_config = load_config("env_config")
+
+    # Set the number of max vehicles
+    env_config.max_num_vehicles = MAX_AGENTS
+
+    # from stable_baselines3.common.vec_env import SubprocVecEnv
+
+    # # Make environment
+    # envs = SubprocVecEnv([lambda: make_env(env_config, MAX_AGENTS) for _ in range(4)])
+    env = make_env(env_config, MAX_AGENTS)
+    import pufferlib.emulation
+    env = pufferlib.emulation.GymnasiumPufferEnv(env, postprocessor_cls=CustomPostprocessor)
+    env.reset()
+    env.step(env.action_space.sample())
+    import pufferlib.vectorization
+    vec = pufferlib.vectorization.Multiprocessing
+    envs = vec(nocturne_creator,env_args=[env_config, MAX_AGENTS], num_envs=4, envs_per_worker=2, env_pool=True)
+    envs.async_reset()
+    obs = envs.recv()[0]
+    actions = [envs.single_action_space.sample() for _ in range(4)]
+    envs.step(actions)
+    envs.step(actions)
+
+    for global_step in range(NUM_STEPS):
+        # Take random action(s) -- you'd obtain this from a policy
+        actions = np.array([envs.action_space.sample() for _ in range(4)])
+
+        # Step
+        obs, rew, done, info = envs.step(actions)
 
-    def get_attr(self, attr_name: str):
-        return getattr(self._env, attr_name)
+        # Log
+        # logging.info(f"step_num: {env.step_num} (global = {global_step}) | done: {done} | rew: {rew}")
 
-    def set_attr(self, attr_name: str):
-        setattr(self._env, attr_name)
+        time.sleep(0.2)