tutorials/4_Stop_Gradient.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `torchopt.stop_gradient` in Meta-Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[<img align=\"left\" src=\"https://colab.research.google.com/assets/colab-badge.svg\">](https://colab.research.google.com/github/metaopt/torchopt/blob/main/tutorials/4_Stop_Gradient.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial, we will illustrate the usage of `torchopt.stop_gradient` with a meta-learning example. We use `torchopt.visual` to help us visualize what is going on in automatic differentiation. Firstly, we define a simple network and the objective function for inner- and outer- optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchopt\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, dim):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Linear(dim, 1, bias=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)\n",
    "\n",
    "\n",
    "loss_fn = F.mse_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the input `x` and output `y`. `y` will be served as the regression target in the following code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "dim = 16\n",
    "\n",
    "x = torch.randn((batch_size, dim))\n",
    "y = torch.zeros((batch_size, 1))\n",
    "net = Net(dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us define the meta-parameter,  we use `MetaSGD` as the inner-loop optimizer and `Adam` as the outer-loop optimizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta_parameter = nn.Parameter(torch.tensor(1.0), requires_grad=True)\n",
    "\n",
    "optim = torchopt.MetaSGD(net, lr=1e-1)\n",
    "meta_optim = torch.optim.Adam([meta_parameter], lr=1e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the inner-loop optimization and visualize the inner-loop forward gradient flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inner loss: 0.3472\n",
      "<graphviz.graphs.Digraph object at 0x7f5a21d70c40>\n"
     ]
    },
    {
     "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.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"226pt\" height=\"335pt\"\n viewBox=\"0.00 0.00 226.00 335.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 331)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-331 222,-331 222,4 -4,4\"/>\n<!-- 140025091550880 -->\n<g id=\"node1\" class=\"node\">\n<title>140025091550880</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"144,-30 67,-30 67,0 144,0 144,-30\"/>\n<text text-anchor=\"middle\" x=\"105.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">inner_loss</text>\n<text text-anchor=\"middle\" x=\"105.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140028156253184 -->\n<g id=\"node2\" class=\"node\">\n<title>140028156253184</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"162,-85 49,-85 49,-66 162,-66 162,-85\"/>\n<text text-anchor=\"middle\" x=\"105.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackward0</text>\n</g>\n<!-- 140028156253184&#45;&gt;140025091550880 -->\n<g id=\"edge7\" class=\"edge\">\n<title>140028156253184&#45;&gt;140025091550880</title>\n<path fill=\"none\" stroke=\"black\" d=\"M105.5,-65.87C105.5,-59.11 105.5,-49.35 105.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"109,-40.11 105.5,-30.11 102,-40.11 109,-40.11\"/>\n</g>\n<!-- 140028156436736 -->\n<g id=\"node3\" class=\"node\">\n<title>140028156436736</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"156,-140 55,-140 55,-121 156,-121 156,-140\"/>\n<text text-anchor=\"middle\" x=\"105.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140028156436736&#45;&gt;140028156253184 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140028156436736&#45;&gt;140028156253184</title>\n<path fill=\"none\" stroke=\"black\" d=\"M105.5,-120.75C105.5,-113.8 105.5,-103.85 105.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"109,-95.09 105.5,-85.09 102,-95.09 109,-95.09\"/>\n</g>\n<!-- 140025091526416 -->\n<g id=\"node4\" class=\"node\">\n<title>140025091526416</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-195 0,-195 0,-176 101,-176 101,-195\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140025091526416&#45;&gt;140028156436736 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140025091526416&#45;&gt;140028156436736</title>\n<path fill=\"none\" stroke=\"black\" d=\"M59.58,-175.75C67.59,-168.03 79.46,-156.6 89.12,-147.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"91.81,-149.55 96.59,-140.09 86.96,-144.51 91.81,-149.55\"/>\n</g>\n<!-- 140028155952000 -->\n<g id=\"node5\" class=\"node\">\n<title>140028155952000</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"96,-261 1,-261 1,-231 96,-231 96,-261\"/>\n<text text-anchor=\"middle\" x=\"48.5\" y=\"-249\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.bias</text>\n<text text-anchor=\"middle\" x=\"48.5\" y=\"-238\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140028155952000&#45;&gt;140025091526416 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140028155952000&#45;&gt;140025091526416</title>\n<path fill=\"none\" stroke=\"black\" d=\"M48.98,-230.84C49.24,-223.21 49.57,-213.7 49.85,-205.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"53.36,-205.38 50.2,-195.27 46.36,-205.14 53.36,-205.38\"/>\n</g>\n<!-- 140025091525408 -->\n<g id=\"node6\" class=\"node\">\n<title>140025091525408</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"200,-195 123,-195 123,-176 200,-176 200,-195\"/>\n<text text-anchor=\"middle\" x=\"161.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091525408&#45;&gt;140028156436736 -->\n<g id=\"edge4\" class=\"edge\">\n<title>140025091525408&#45;&gt;140028156436736</title>\n<path fill=\"none\" stroke=\"black\" d=\"M152.5,-175.98C144.31,-168.23 131.99,-156.58 122.03,-147.14\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"124.33,-144.5 114.66,-140.17 119.52,-149.59 124.33,-144.5\"/>\n</g>\n<!-- 140025091526224 -->\n<g id=\"node7\" class=\"node\">\n<title>140025091526224</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"215,-255.5 114,-255.5 114,-236.5 215,-236.5 215,-255.5\"/>\n<text text-anchor=\"middle\" x=\"164.5\" y=\"-243.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140025091526224&#45;&gt;140025091525408 -->\n<g id=\"edge5\" class=\"edge\">\n<title>140025091526224&#45;&gt;140025091525408</title>\n<path fill=\"none\" stroke=\"black\" d=\"M164.06,-236.37C163.64,-228.25 163,-215.81 162.47,-205.39\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"165.95,-204.97 161.94,-195.17 158.96,-205.33 165.95,-204.97\"/>\n</g>\n<!-- 140028155952880 -->\n<g id=\"node8\" class=\"node\">\n<title>140028155952880</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"218,-327 111,-327 111,-297 218,-297 218,-327\"/>\n<text text-anchor=\"middle\" x=\"164.5\" y=\"-315\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.weight</text>\n<text text-anchor=\"middle\" x=\"164.5\" y=\"-304\" font-family=\"monospace\" font-size=\"10.00\">(1, 16)</text>\n</g>\n<!-- 140028155952880&#45;&gt;140025091526224 -->\n<g id=\"edge6\" class=\"edge\">\n<title>140028155952880&#45;&gt;140025091526224</title>\n<path fill=\"none\" stroke=\"black\" d=\"M164.5,-296.8C164.5,-287.7 164.5,-275.79 164.5,-265.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"168,-265.84 164.5,-255.84 161,-265.84 168,-265.84\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "init_net_state = torchopt.extract_state_dict(net, enable_visual=True, visual_prefix='step0.')\n",
    "\n",
    "# inner loss\n",
    "inner_loss = loss_fn(net(x), y)\n",
    "\n",
    "print(f'inner loss: {inner_loss:.4f}')\n",
    "display(torchopt.visual.make_dot(inner_loss, params=(init_net_state, {'inner_loss': inner_loss})))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Conduct inner-loop optimization with `MetaSGD`, here the meta-parameter is served as a factor controlling the scale of inner-loop loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inner-step optimization\n",
    "loss = inner_loss * meta_parameter\n",
    "optim.step(loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We compute the outer loss and draw the full computation graph of the first bi-level process. In this graph, three main parts are included.\n",
    "\n",
    "- Inner-loop: forward process and inner-loss calculation\n",
    "- Inner-loop optimization: `MetaSGD` optimization step given inner-loss\n",
    "- Outer-loop: forward process and outer-loss calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "outer loss: 0.2039\n",
      "<graphviz.graphs.Digraph object at 0x7f5a21d70730>\n"
     ]
    },
    {
     "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.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"421pt\" height=\"874pt\"\n viewBox=\"0.00 0.00 421.00 874.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 870)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-870 417,-870 417,4 -4,4\"/>\n<!-- 140027829238416 -->\n<g id=\"node1\" class=\"node\">\n<title>140027829238416</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"189,-30 112,-30 112,0 189,0 189,-30\"/>\n<text text-anchor=\"middle\" x=\"150.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">outer_loss</text>\n<text text-anchor=\"middle\" x=\"150.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140025091525072 -->\n<g id=\"node2\" class=\"node\">\n<title>140025091525072</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"207,-85 94,-85 94,-66 207,-66 207,-85\"/>\n<text text-anchor=\"middle\" x=\"150.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackward0</text>\n</g>\n<!-- 140025091525072&#45;&gt;140027829238416 -->\n<g id=\"edge26\" class=\"edge\">\n<title>140025091525072&#45;&gt;140027829238416</title>\n<path fill=\"none\" stroke=\"black\" d=\"M150.5,-65.87C150.5,-59.11 150.5,-49.35 150.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"154,-40.11 150.5,-30.11 147,-40.11 154,-40.11\"/>\n</g>\n<!-- 140025091525216 -->\n<g id=\"node3\" class=\"node\">\n<title>140025091525216</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"201,-140 100,-140 100,-121 201,-121 201,-140\"/>\n<text text-anchor=\"middle\" x=\"150.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140025091525216&#45;&gt;140025091525072 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140025091525216&#45;&gt;140025091525072</title>\n<path fill=\"none\" stroke=\"black\" d=\"M150.5,-120.75C150.5,-113.8 150.5,-103.85 150.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"154,-95.09 150.5,-85.09 147,-95.09 154,-95.09\"/>\n</g>\n<!-- 140025091526128 -->\n<g id=\"node4\" class=\"node\">\n<title>140025091526128</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"154,-404 59,-404 59,-363 154,-363 154,-404\"/>\n<text text-anchor=\"middle\" x=\"106.5\" y=\"-392\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n<text text-anchor=\"middle\" x=\"106.5\" y=\"-381\" font-family=\"monospace\" font-size=\"10.00\">step1.fc.bias</text>\n<text text-anchor=\"middle\" x=\"106.5\" y=\"-370\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140025091526128&#45;&gt;140025091525216 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140025091526128&#45;&gt;140025091525216</title>\n<path fill=\"none\" stroke=\"black\" d=\"M109.95,-362.8C118.22,-315.61 139.09,-196.57 147.2,-150.3\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"150.66,-150.88 148.94,-140.42 143.76,-149.67 150.66,-150.88\"/>\n</g>\n<!-- 140025091526416 -->\n<g id=\"node5\" class=\"node\">\n<title>140025091526416</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-734 0,-734 0,-715 101,-715 101,-734\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-722\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140025091526416&#45;&gt;140025091526128 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140025091526416&#45;&gt;140025091526128</title>\n<path fill=\"none\" stroke=\"black\" d=\"M47.8,-714.92C42.34,-696.77 30.5,-653.07 30.5,-615.5 30.5,-615.5 30.5,-615.5 30.5,-503.5 30.5,-473.63 36.82,-465.42 52.5,-440 58.83,-429.73 67.36,-419.85 75.76,-411.33\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"78.33,-413.71 83.06,-404.23 73.45,-408.7 78.33,-413.71\"/>\n</g>\n<!-- 140028156436736 -->\n<g id=\"node14\" class=\"node\">\n<title>140028156436736</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"225,-679 124,-679 124,-660 225,-660 225,-679\"/>\n<text text-anchor=\"middle\" x=\"174.5\" y=\"-667\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140025091526416&#45;&gt;140028156436736 -->\n<g id=\"edge13\" class=\"edge\">\n<title>140025091526416&#45;&gt;140028156436736</title>\n<path fill=\"none\" stroke=\"black\" d=\"M70.42,-714.98C90.55,-706.38 121.88,-692.99 144.91,-683.15\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"146.4,-686.32 154.21,-679.17 143.64,-679.88 146.4,-686.32\"/>\n</g>\n<!-- 140028155952000 -->\n<g id=\"node6\" class=\"node\">\n<title>140028155952000</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"98,-800 3,-800 3,-770 98,-770 98,-800\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-788\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.bias</text>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-777\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140028155952000&#45;&gt;140025091526416 -->\n<g id=\"edge4\" class=\"edge\">\n<title>140028155952000&#45;&gt;140025091526416</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-769.84C50.5,-762.21 50.5,-752.7 50.5,-744.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-744.27 50.5,-734.27 47,-744.27 54,-744.27\"/>\n</g>\n<!-- 140025091524976 -->\n<g id=\"node7\" class=\"node\">\n<title>140025091524976</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"151,-459 62,-459 62,-440 151,-440 151,-459\"/>\n<text text-anchor=\"middle\" x=\"106.5\" y=\"-447\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140025091524976&#45;&gt;140025091526128 -->\n<g id=\"edge5\" class=\"edge\">\n<title>140025091524976&#45;&gt;140025091526128</title>\n<path fill=\"none\" stroke=\"black\" d=\"M106.5,-439.87C106.5,-433.22 106.5,-423.63 106.5,-414.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"110,-414.01 106.5,-404.01 103,-414.01 110,-414.01\"/>\n</g>\n<!-- 140025091526560 -->\n<g id=\"node8\" class=\"node\">\n<title>140025091526560</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"154,-514 59,-514 59,-495 154,-495 154,-514\"/>\n<text text-anchor=\"middle\" x=\"106.5\" y=\"-502\" font-family=\"monospace\" font-size=\"10.00\">ViewBackward0</text>\n</g>\n<!-- 140025091526560&#45;&gt;140025091524976 -->\n<g id=\"edge6\" class=\"edge\">\n<title>140025091526560&#45;&gt;140025091524976</title>\n<path fill=\"none\" stroke=\"black\" d=\"M106.5,-494.75C106.5,-487.8 106.5,-477.85 106.5,-469.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"110,-469.09 106.5,-459.09 103,-469.09 110,-469.09\"/>\n</g>\n<!-- 140025091525456 -->\n<g id=\"node9\" class=\"node\">\n<title>140025091525456</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"152,-569 63,-569 63,-550 152,-550 152,-569\"/>\n<text text-anchor=\"middle\" x=\"107.5\" y=\"-557\" font-family=\"monospace\" font-size=\"10.00\">SumBackward1</text>\n</g>\n<!-- 140025091525456&#45;&gt;140025091526560 -->\n<g id=\"edge7\" class=\"edge\">\n<title>140025091525456&#45;&gt;140025091526560</title>\n<path fill=\"none\" stroke=\"black\" d=\"M107.33,-549.75C107.2,-542.8 107.02,-532.85 106.85,-524.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"110.35,-524.02 106.66,-514.09 103.35,-524.15 110.35,-524.02\"/>\n</g>\n<!-- 140025091524112 -->\n<g id=\"node10\" class=\"node\">\n<title>140025091524112</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"255,-624 94,-624 94,-605 255,-605 255,-624\"/>\n<text text-anchor=\"middle\" x=\"174.5\" y=\"-612\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackwardBackward0</text>\n</g>\n<!-- 140025091524112&#45;&gt;140025091525456 -->\n<g id=\"edge8\" class=\"edge\">\n<title>140025091524112&#45;&gt;140025091525456</title>\n<path fill=\"none\" stroke=\"black\" d=\"M163.74,-604.98C153.73,-597.07 138.61,-585.11 126.57,-575.58\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"128.47,-572.63 118.46,-569.17 124.13,-578.12 128.47,-572.63\"/>\n</g>\n<!-- 140024973742672 -->\n<g id=\"node24\" class=\"node\">\n<title>140024973742672</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"249,-569 172,-569 172,-550 249,-550 249,-569\"/>\n<text text-anchor=\"middle\" x=\"210.5\" y=\"-557\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091524112&#45;&gt;140024973742672 -->\n<g id=\"edge25\" class=\"edge\">\n<title>140025091524112&#45;&gt;140024973742672</title>\n<path fill=\"none\" stroke=\"black\" d=\"M180.44,-604.75C185.48,-597.34 192.84,-586.5 199.01,-577.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"201.94,-579.33 204.67,-569.09 196.15,-575.39 201.94,-579.33\"/>\n</g>\n<!-- 140024973742288 -->\n<g id=\"node11\" class=\"node\">\n<title>140024973742288</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"404,-679 315,-679 315,-660 404,-660 404,-679\"/>\n<text text-anchor=\"middle\" x=\"359.5\" y=\"-667\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140024973742288&#45;&gt;140025091524112 -->\n<g id=\"edge9\" class=\"edge\">\n<title>140024973742288&#45;&gt;140025091524112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M329.78,-659.98C298.45,-651.01 248.93,-636.82 214.26,-626.89\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"215.03,-623.47 204.46,-624.08 213.11,-630.2 215.03,-623.47\"/>\n</g>\n<!-- 140024973742384 -->\n<g id=\"node12\" class=\"node\">\n<title>140024973742384</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"413,-734 312,-734 312,-715 413,-715 413,-734\"/>\n<text text-anchor=\"middle\" x=\"362.5\" y=\"-722\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140024973742384&#45;&gt;140024973742288 -->\n<g id=\"edge10\" class=\"edge\">\n<title>140024973742384&#45;&gt;140024973742288</title>\n<path fill=\"none\" stroke=\"black\" d=\"M362,-714.75C361.61,-707.8 361.05,-697.85 360.55,-689.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"364.05,-688.88 359.99,-679.09 357.06,-689.27 364.05,-688.88\"/>\n</g>\n<!-- 140025091549440 -->\n<g id=\"node13\" class=\"node\">\n<title>140025091549440</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"413,-800 312,-800 312,-770 413,-770 413,-800\"/>\n<text text-anchor=\"middle\" x=\"362.5\" y=\"-788\" font-family=\"monospace\" font-size=\"10.00\">meta_parameter</text>\n<text text-anchor=\"middle\" x=\"362.5\" y=\"-777\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140025091549440&#45;&gt;140024973742384 -->\n<g id=\"edge11\" class=\"edge\">\n<title>140025091549440&#45;&gt;140024973742384</title>\n<path fill=\"none\" stroke=\"black\" d=\"M362.5,-769.84C362.5,-762.21 362.5,-752.7 362.5,-744.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"366,-744.27 362.5,-734.27 359,-744.27 366,-744.27\"/>\n</g>\n<!-- 140028156436736&#45;&gt;140025091524112 -->\n<g id=\"edge12\" class=\"edge\">\n<title>140028156436736&#45;&gt;140025091524112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M174.5,-659.75C174.5,-652.8 174.5,-642.85 174.5,-634.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"178,-634.09 174.5,-624.09 171,-634.09 178,-634.09\"/>\n</g>\n<!-- 140025091525408 -->\n<g id=\"node15\" class=\"node\">\n<title>140025091525408</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"235,-734 158,-734 158,-715 235,-715 235,-734\"/>\n<text text-anchor=\"middle\" x=\"196.5\" y=\"-722\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091525408&#45;&gt;140028156436736 -->\n<g id=\"edge14\" class=\"edge\">\n<title>140025091525408&#45;&gt;140028156436736</title>\n<path fill=\"none\" stroke=\"black\" d=\"M192.87,-714.75C189.92,-707.65 185.67,-697.4 181.99,-688.56\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"185.13,-686.98 178.07,-679.09 178.67,-689.67 185.13,-686.98\"/>\n</g>\n<!-- 140025091526224 -->\n<g id=\"node16\" class=\"node\">\n<title>140025091526224</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"285,-794.5 184,-794.5 184,-775.5 285,-775.5 285,-794.5\"/>\n<text text-anchor=\"middle\" x=\"234.5\" y=\"-782.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140025091526224&#45;&gt;140025091525408 -->\n<g id=\"edge15\" class=\"edge\">\n<title>140025091526224&#45;&gt;140025091525408</title>\n<path fill=\"none\" stroke=\"black\" d=\"M228.89,-775.37C223.34,-766.81 214.67,-753.47 207.7,-742.74\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"210.51,-740.65 202.13,-734.17 204.64,-744.46 210.51,-740.65\"/>\n</g>\n<!-- 140025091524928 -->\n<g id=\"node19\" class=\"node\">\n<title>140025091524928</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"264,-272 157,-272 157,-231 264,-231 264,-272\"/>\n<text text-anchor=\"middle\" x=\"210.5\" y=\"-260\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n<text text-anchor=\"middle\" x=\"210.5\" y=\"-249\" font-family=\"monospace\" font-size=\"10.00\">step1.fc.weight</text>\n<text text-anchor=\"middle\" x=\"210.5\" y=\"-238\" font-family=\"monospace\" font-size=\"10.00\">(1, 16)</text>\n</g>\n<!-- 140025091526224&#45;&gt;140025091524928 -->\n<g id=\"edge19\" class=\"edge\">\n<title>140025091526224&#45;&gt;140025091524928</title>\n<path fill=\"none\" stroke=\"black\" d=\"M241.11,-775.43C254.48,-757.2 283.5,-712.78 283.5,-670.5 283.5,-670.5 283.5,-670.5 283.5,-382.5 283.5,-348.22 280.31,-337.88 263.5,-308 257.64,-297.58 249.38,-287.65 241.12,-279.14\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"243.5,-276.58 233.92,-272.05 238.59,-281.56 243.5,-276.58\"/>\n</g>\n<!-- 140028155952880 -->\n<g id=\"node17\" class=\"node\">\n<title>140028155952880</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"288,-866 181,-866 181,-836 288,-836 288,-866\"/>\n<text text-anchor=\"middle\" x=\"234.5\" y=\"-854\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.weight</text>\n<text text-anchor=\"middle\" x=\"234.5\" y=\"-843\" font-family=\"monospace\" font-size=\"10.00\">(1, 16)</text>\n</g>\n<!-- 140028155952880&#45;&gt;140025091526224 -->\n<g id=\"edge16\" class=\"edge\">\n<title>140028155952880&#45;&gt;140025091526224</title>\n<path fill=\"none\" stroke=\"black\" d=\"M234.5,-835.8C234.5,-826.7 234.5,-814.79 234.5,-804.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"238,-804.84 234.5,-794.84 231,-804.84 238,-804.84\"/>\n</g>\n<!-- 140025091524448 -->\n<g id=\"node18\" class=\"node\">\n<title>140025091524448</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"230,-195 153,-195 153,-176 230,-176 230,-195\"/>\n<text text-anchor=\"middle\" x=\"191.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091524448&#45;&gt;140025091525216 -->\n<g id=\"edge17\" class=\"edge\">\n<title>140025091524448&#45;&gt;140025091525216</title>\n<path fill=\"none\" stroke=\"black\" d=\"M184.73,-175.75C178.94,-168.26 170.44,-157.28 163.36,-148.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"166.03,-145.86 157.14,-140.09 160.49,-150.14 166.03,-145.86\"/>\n</g>\n<!-- 140025091524928&#45;&gt;140025091524448 -->\n<g id=\"edge18\" class=\"edge\">\n<title>140025091524928&#45;&gt;140025091524448</title>\n<path fill=\"none\" stroke=\"black\" d=\"M204.7,-230.95C202.24,-222.67 199.39,-213.07 196.97,-204.92\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"200.3,-203.85 194.1,-195.26 193.59,-205.84 200.3,-203.85\"/>\n</g>\n<!-- 140025091525600 -->\n<g id=\"node20\" class=\"node\">\n<title>140025091525600</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"255,-327 166,-327 166,-308 255,-308 255,-327\"/>\n<text text-anchor=\"middle\" x=\"210.5\" y=\"-315\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140025091525600&#45;&gt;140025091524928 -->\n<g id=\"edge20\" class=\"edge\">\n<title>140025091525600&#45;&gt;140025091524928</title>\n<path fill=\"none\" stroke=\"black\" d=\"M210.5,-307.87C210.5,-301.22 210.5,-291.63 210.5,-282.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"214,-282.01 210.5,-272.01 207,-282.01 214,-282.01\"/>\n</g>\n<!-- 140024973742144 -->\n<g id=\"node21\" class=\"node\">\n<title>140024973742144</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"250,-393 173,-393 173,-374 250,-374 250,-393\"/>\n<text text-anchor=\"middle\" x=\"211.5\" y=\"-381\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973742144&#45;&gt;140025091525600 -->\n<g id=\"edge21\" class=\"edge\">\n<title>140024973742144&#45;&gt;140025091525600</title>\n<path fill=\"none\" stroke=\"black\" d=\"M211.37,-373.87C211.22,-364.66 210.99,-349.79 210.8,-337.77\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"214.3,-337.35 210.64,-327.41 207.3,-337.46 214.3,-337.35\"/>\n</g>\n<!-- 140024973742576 -->\n<g id=\"node22\" class=\"node\">\n<title>140024973742576</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"251,-459 174,-459 174,-440 251,-440 251,-459\"/>\n<text text-anchor=\"middle\" x=\"212.5\" y=\"-447\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973742576&#45;&gt;140024973742144 -->\n<g id=\"edge22\" class=\"edge\">\n<title>140024973742576&#45;&gt;140024973742144</title>\n<path fill=\"none\" stroke=\"black\" d=\"M212.37,-439.87C212.22,-430.66 211.99,-415.79 211.8,-403.77\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"215.3,-403.35 211.64,-393.41 208.3,-403.46 215.3,-403.35\"/>\n</g>\n<!-- 140024973742480 -->\n<g id=\"node23\" class=\"node\">\n<title>140024973742480</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"255,-514 172,-514 172,-495 255,-495 255,-514\"/>\n<text text-anchor=\"middle\" x=\"213.5\" y=\"-502\" font-family=\"monospace\" font-size=\"10.00\">MmBackward0</text>\n</g>\n<!-- 140024973742480&#45;&gt;140024973742576 -->\n<g id=\"edge23\" class=\"edge\">\n<title>140024973742480&#45;&gt;140024973742576</title>\n<path fill=\"none\" stroke=\"black\" d=\"M213.33,-494.75C213.2,-487.8 213.02,-477.85 212.85,-469.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"216.35,-469.02 212.66,-459.09 209.35,-469.15 216.35,-469.02\"/>\n</g>\n<!-- 140024973742672&#45;&gt;140024973742480 -->\n<g id=\"edge24\" class=\"edge\">\n<title>140024973742672&#45;&gt;140024973742480</title>\n<path fill=\"none\" stroke=\"black\" d=\"M211,-549.75C211.39,-542.8 211.95,-532.85 212.45,-524.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"215.94,-524.27 213.01,-514.09 208.95,-523.88 215.94,-524.27\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Extract `state_dict`` for updated network\n",
    "one_step_net_state = torchopt.extract_state_dict(net, enable_visual=True, visual_prefix='step1.')\n",
    "one_step_optim_state = torchopt.extract_state_dict(optim)\n",
    "\n",
    "# Calculate outer loss\n",
    "outer_loss = loss_fn(net(x), y)\n",
    "print(f'outer loss: {outer_loss:.4f}')\n",
    "display(\n",
    "    torchopt.visual.make_dot(\n",
    "        outer_loss,\n",
    "        params=(\n",
    "            init_net_state,\n",
    "            one_step_net_state,\n",
    "            {'meta_parameter': meta_parameter, 'outer_loss': outer_loss},\n",
    "        ),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we backward the loss to conduct outer-loop meta-optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "meta_parameter.grad = tensor(-0.1205)\n",
      "meta_parameter = Parameter containing:\n",
      "tensor(1.1000, requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "meta_optim.zero_grad()\n",
    "outer_loss.backward()\n",
    "print(f'meta_parameter.grad = {meta_parameter.grad!r}')\n",
    "meta_optim.step()\n",
    "print(f'meta_parameter = {meta_parameter!r}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have already conducted one bi-level optimization and optimize our meta-parameters. When you want to conduct the second bi-level optimization, you need to be careful whether you need to use the `stop_gradient` function. For example, if your new inner-loop parameters directly inherits previous inner-loop parameters (which is a common strategy in many meta-learning algorithms like Meta-Gradient Reinforcement Learning (MGRL) ([arXiv:1805.09801](https://arxiv.org/abs/1805.09801))), you might need `stop_gradient` function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, the backpropagation only frees saved tensors (often used as auxiliary data for computing the gradient) but the computation graph remains. Once the outer iteration is finished, if you want to use any intermediate network parameters produced by the inner loop for the next bi-level iteration, you should detach them from the computation graph.\n",
    "\n",
    "There are two main reasons:\n",
    "\n",
    "- The network parameters are still connected to the previous computation graph (`.grad_fn` is not `None`). If later the gradient backpropagate to these parameters, the PyTorch backward engine will try to backpropagate through the previous computation graph. This will raise a `RuntimeError`: Trying to backward through the graph a second time...\n",
    "- If we do not detach the computation graph, the computation graph connected to these parameters can not be freed by GC (Garbage Collector) until these parameters are collected by GC."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us see what will happen if we do not use the `stop_gradient` function before we conduct the second bi-level process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<graphviz.graphs.Digraph object at 0x7f5ac5072280>\n"
     ]
    },
    {
     "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.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"413pt\" height=\"1369pt\"\n viewBox=\"0.00 0.00 413.00 1369.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 1365)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-1365 409,-1365 409,4 -4,4\"/>\n<!-- 140024973755152 -->\n<g id=\"node1\" class=\"node\">\n<title>140024973755152</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"221,-30 144,-30 144,0 221,0 221,-30\"/>\n<text text-anchor=\"middle\" x=\"182.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">outer_loss</text>\n<text text-anchor=\"middle\" x=\"182.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140027829363232 -->\n<g id=\"node2\" class=\"node\">\n<title>140027829363232</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"239,-85 126,-85 126,-66 239,-66 239,-85\"/>\n<text text-anchor=\"middle\" x=\"182.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackward0</text>\n</g>\n<!-- 140027829363232&#45;&gt;140024973755152 -->\n<g id=\"edge44\" class=\"edge\">\n<title>140027829363232&#45;&gt;140024973755152</title>\n<path fill=\"none\" stroke=\"black\" d=\"M182.5,-65.87C182.5,-59.11 182.5,-49.35 182.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"186,-40.11 182.5,-30.11 179,-40.11 186,-40.11\"/>\n</g>\n<!-- 140027829363616 -->\n<g id=\"node3\" class=\"node\">\n<title>140027829363616</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"233,-140 132,-140 132,-121 233,-121 233,-140\"/>\n<text text-anchor=\"middle\" x=\"182.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140027829363616&#45;&gt;140027829363232 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140027829363616&#45;&gt;140027829363232</title>\n<path fill=\"none\" stroke=\"black\" d=\"M182.5,-120.75C182.5,-113.8 182.5,-103.85 182.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"186,-95.09 182.5,-85.09 179,-95.09 186,-95.09\"/>\n</g>\n<!-- 140027829366544 -->\n<g id=\"node4\" class=\"node\">\n<title>140027829366544</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"281,-360 192,-360 192,-341 281,-341 281,-360\"/>\n<text text-anchor=\"middle\" x=\"236.5\" y=\"-348\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n</g>\n<!-- 140027829366544&#45;&gt;140027829363616 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140027829366544&#45;&gt;140027829363616</title>\n<path fill=\"none\" stroke=\"black\" d=\"M234.74,-340.77C229.54,-315.26 213.55,-238.64 196.5,-176 194.16,-167.4 191.25,-157.98 188.68,-150.02\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"191.94,-148.71 185.49,-140.31 185.29,-150.9 191.94,-148.71\"/>\n</g>\n<!-- 140025091526128 -->\n<g id=\"node5\" class=\"node\">\n<title>140025091526128</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"360,-844 265,-844 265,-803 360,-803 360,-844\"/>\n<text text-anchor=\"middle\" x=\"312.5\" y=\"-832\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n<text text-anchor=\"middle\" x=\"312.5\" y=\"-821\" font-family=\"monospace\" font-size=\"10.00\">step1.fc.bias</text>\n<text text-anchor=\"middle\" x=\"312.5\" y=\"-810\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140025091526128&#45;&gt;140027829366544 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140025091526128&#45;&gt;140027829366544</title>\n<path fill=\"none\" stroke=\"black\" d=\"M312.5,-802.88C312.5,-775.65 312.5,-724.83 312.5,-681.5 312.5,-681.5 312.5,-681.5 312.5,-459.5 312.5,-429.48 306.69,-420.6 289.5,-396 281.33,-384.31 269.33,-373.88 258.82,-366.07\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"260.64,-363.07 250.45,-360.15 256.6,-368.79 260.64,-363.07\"/>\n</g>\n<!-- 140025091725152 -->\n<g id=\"node24\" class=\"node\">\n<title>140025091725152</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"226,-635 125,-635 125,-616 226,-616 226,-635\"/>\n<text text-anchor=\"middle\" x=\"175.5\" y=\"-623\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140025091526128&#45;&gt;140025091725152 -->\n<g id=\"edge25\" class=\"edge\">\n<title>140025091526128&#45;&gt;140025091725152</title>\n<path fill=\"none\" stroke=\"black\" d=\"M302.16,-802.98C286.03,-773.27 253.12,-715.44 218.5,-671 210.6,-660.85 200.68,-650.46 192.36,-642.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"194.63,-639.6 184.99,-635.18 189.77,-644.64 194.63,-639.6\"/>\n</g>\n<!-- 140025091526416 -->\n<g id=\"node6\" class=\"node\">\n<title>140025091526416</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"405,-1229 304,-1229 304,-1210 405,-1210 405,-1229\"/>\n<text text-anchor=\"middle\" x=\"354.5\" y=\"-1217\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140025091526416&#45;&gt;140025091526128 -->\n<g id=\"edge4\" class=\"edge\">\n<title>140025091526416&#45;&gt;140025091526128</title>\n<path fill=\"none\" stroke=\"black\" d=\"M360.06,-1209.94C371.05,-1192.18 394.5,-1149.69 394.5,-1110.5 394.5,-1110.5 394.5,-1110.5 394.5,-943.5 394.5,-906.76 367.67,-873.11 344.7,-851.04\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"346.82,-848.23 337.1,-844.01 342.07,-853.37 346.82,-848.23\"/>\n</g>\n<!-- 140028156436736 -->\n<g id=\"node15\" class=\"node\">\n<title>140028156436736</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"366,-1174 265,-1174 265,-1155 366,-1155 366,-1174\"/>\n<text text-anchor=\"middle\" x=\"315.5\" y=\"-1162\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140025091526416&#45;&gt;140028156436736 -->\n<g id=\"edge14\" class=\"edge\">\n<title>140025091526416&#45;&gt;140028156436736</title>\n<path fill=\"none\" stroke=\"black\" d=\"M348.06,-1209.75C342.61,-1202.34 334.64,-1191.5 327.94,-1182.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"330.57,-1180.07 321.82,-1174.09 324.93,-1184.22 330.57,-1180.07\"/>\n</g>\n<!-- 140028155952000 -->\n<g id=\"node7\" class=\"node\">\n<title>140028155952000</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"402,-1295 307,-1295 307,-1265 402,-1265 402,-1295\"/>\n<text text-anchor=\"middle\" x=\"354.5\" y=\"-1283\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.bias</text>\n<text text-anchor=\"middle\" x=\"354.5\" y=\"-1272\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140028155952000&#45;&gt;140025091526416 -->\n<g id=\"edge5\" class=\"edge\">\n<title>140028155952000&#45;&gt;140025091526416</title>\n<path fill=\"none\" stroke=\"black\" d=\"M354.5,-1264.84C354.5,-1257.21 354.5,-1247.7 354.5,-1239.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"358,-1239.27 354.5,-1229.27 351,-1239.27 358,-1239.27\"/>\n</g>\n<!-- 140025091524976 -->\n<g id=\"node8\" class=\"node\">\n<title>140025091524976</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"357,-954 268,-954 268,-935 357,-935 357,-954\"/>\n<text text-anchor=\"middle\" x=\"312.5\" y=\"-942\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140025091524976&#45;&gt;140025091526128 -->\n<g id=\"edge6\" class=\"edge\">\n<title>140025091524976&#45;&gt;140025091526128</title>\n<path fill=\"none\" stroke=\"black\" d=\"M312.5,-934.94C312.5,-918.36 312.5,-881.15 312.5,-854.5\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"316,-854.17 312.5,-844.17 309,-854.17 316,-854.17\"/>\n</g>\n<!-- 140025091526560 -->\n<g id=\"node9\" class=\"node\">\n<title>140025091526560</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"360,-1009 265,-1009 265,-990 360,-990 360,-1009\"/>\n<text text-anchor=\"middle\" x=\"312.5\" y=\"-997\" font-family=\"monospace\" font-size=\"10.00\">ViewBackward0</text>\n</g>\n<!-- 140025091526560&#45;&gt;140025091524976 -->\n<g id=\"edge7\" class=\"edge\">\n<title>140025091526560&#45;&gt;140025091524976</title>\n<path fill=\"none\" stroke=\"black\" d=\"M312.5,-989.75C312.5,-982.8 312.5,-972.85 312.5,-964.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"316,-964.09 312.5,-954.09 309,-964.09 316,-964.09\"/>\n</g>\n<!-- 140025091525456 -->\n<g id=\"node10\" class=\"node\">\n<title>140025091525456</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"355,-1064 266,-1064 266,-1045 355,-1045 355,-1064\"/>\n<text text-anchor=\"middle\" x=\"310.5\" y=\"-1052\" font-family=\"monospace\" font-size=\"10.00\">SumBackward1</text>\n</g>\n<!-- 140025091525456&#45;&gt;140025091526560 -->\n<g id=\"edge8\" class=\"edge\">\n<title>140025091525456&#45;&gt;140025091526560</title>\n<path fill=\"none\" stroke=\"black\" d=\"M310.83,-1044.75C311.09,-1037.8 311.47,-1027.85 311.8,-1019.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"315.3,-1019.21 312.18,-1009.09 308.3,-1018.95 315.3,-1019.21\"/>\n</g>\n<!-- 140025091524112 -->\n<g id=\"node11\" class=\"node\">\n<title>140025091524112</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"343,-1119 182,-1119 182,-1100 343,-1100 343,-1119\"/>\n<text text-anchor=\"middle\" x=\"262.5\" y=\"-1107\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackwardBackward0</text>\n</g>\n<!-- 140025091524112&#45;&gt;140025091525456 -->\n<g id=\"edge9\" class=\"edge\">\n<title>140025091524112&#45;&gt;140025091525456</title>\n<path fill=\"none\" stroke=\"black\" d=\"M270.43,-1099.75C277.35,-1092.11 287.57,-1080.82 295.95,-1071.56\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"298.6,-1073.85 302.72,-1064.09 293.41,-1069.15 298.6,-1073.85\"/>\n</g>\n<!-- 140024973742672 -->\n<g id=\"node31\" class=\"node\">\n<title>140024973742672</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"245,-1064 168,-1064 168,-1045 245,-1045 245,-1064\"/>\n<text text-anchor=\"middle\" x=\"206.5\" y=\"-1052\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091524112&#45;&gt;140024973742672 -->\n<g id=\"edge34\" class=\"edge\">\n<title>140025091524112&#45;&gt;140024973742672</title>\n<path fill=\"none\" stroke=\"black\" d=\"M253.5,-1099.98C245.31,-1092.23 232.99,-1080.58 223.03,-1071.14\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"225.33,-1068.5 215.66,-1064.17 220.52,-1073.59 225.33,-1068.5\"/>\n</g>\n<!-- 140024973742288 -->\n<g id=\"node12\" class=\"node\">\n<title>140024973742288</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"247,-1174 158,-1174 158,-1155 247,-1155 247,-1174\"/>\n<text text-anchor=\"middle\" x=\"202.5\" y=\"-1162\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140024973742288&#45;&gt;140025091524112 -->\n<g id=\"edge10\" class=\"edge\">\n<title>140024973742288&#45;&gt;140025091524112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M212.14,-1154.98C221.01,-1147.15 234.37,-1135.34 245.11,-1125.86\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"247.51,-1128.41 252.68,-1119.17 242.87,-1123.17 247.51,-1128.41\"/>\n</g>\n<!-- 140024973742384 -->\n<g id=\"node13\" class=\"node\">\n<title>140024973742384</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"101,-1229 0,-1229 0,-1210 101,-1210 101,-1229\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-1217\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140024973742384&#45;&gt;140024973742288 -->\n<g id=\"edge11\" class=\"edge\">\n<title>140024973742384&#45;&gt;140024973742288</title>\n<path fill=\"none\" stroke=\"black\" d=\"M74.92,-1209.98C100.22,-1201.16 139.95,-1187.31 168.36,-1177.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"169.6,-1180.68 177.89,-1174.08 167.29,-1174.07 169.6,-1180.68\"/>\n</g>\n<!-- 140025091726064 -->\n<g id=\"node23\" class=\"node\">\n<title>140025091726064</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"100,-1119 11,-1119 11,-1100 100,-1100 100,-1119\"/>\n<text text-anchor=\"middle\" x=\"55.5\" y=\"-1107\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140024973742384&#45;&gt;140025091726064 -->\n<g id=\"edge23\" class=\"edge\">\n<title>140024973742384&#45;&gt;140025091726064</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.91,-1209.66C51.72,-1192.17 53.54,-1152.8 54.63,-1129.27\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"58.13,-1129.31 55.1,-1119.16 51.14,-1128.99 58.13,-1129.31\"/>\n</g>\n<!-- 140025091549440 -->\n<g id=\"node14\" class=\"node\">\n<title>140025091549440</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"101,-1295 0,-1295 0,-1265 101,-1265 101,-1295\"/>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-1283\" font-family=\"monospace\" font-size=\"10.00\">meta_parameter</text>\n<text text-anchor=\"middle\" x=\"50.5\" y=\"-1272\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140025091549440&#45;&gt;140024973742384 -->\n<g id=\"edge12\" class=\"edge\">\n<title>140025091549440&#45;&gt;140024973742384</title>\n<path fill=\"none\" stroke=\"black\" d=\"M50.5,-1264.84C50.5,-1257.21 50.5,-1247.7 50.5,-1239.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"54,-1239.27 50.5,-1229.27 47,-1239.27 54,-1239.27\"/>\n</g>\n<!-- 140028156436736&#45;&gt;140025091524112 -->\n<g id=\"edge13\" class=\"edge\">\n<title>140028156436736&#45;&gt;140025091524112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M306.75,-1154.75C299.03,-1147.03 287.6,-1135.6 278.28,-1126.28\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"280.64,-1123.69 271.09,-1119.09 275.69,-1128.64 280.64,-1123.69\"/>\n</g>\n<!-- 140025091525408 -->\n<g id=\"node16\" class=\"node\">\n<title>140025091525408</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"280,-1229 203,-1229 203,-1210 280,-1210 280,-1229\"/>\n<text text-anchor=\"middle\" x=\"241.5\" y=\"-1217\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091525408&#45;&gt;140028156436736 -->\n<g id=\"edge15\" class=\"edge\">\n<title>140025091525408&#45;&gt;140028156436736</title>\n<path fill=\"none\" stroke=\"black\" d=\"M253.39,-1209.98C264.65,-1201.92 281.79,-1189.65 295.21,-1180.03\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"297.3,-1182.84 303.39,-1174.17 293.23,-1177.15 297.3,-1182.84\"/>\n</g>\n<!-- 140025091526224 -->\n<g id=\"node17\" class=\"node\">\n<title>140025091526224</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"274,-1289.5 173,-1289.5 173,-1270.5 274,-1270.5 274,-1289.5\"/>\n<text text-anchor=\"middle\" x=\"223.5\" y=\"-1277.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140025091526224&#45;&gt;140025091525408 -->\n<g id=\"edge16\" class=\"edge\">\n<title>140025091526224&#45;&gt;140025091525408</title>\n<path fill=\"none\" stroke=\"black\" d=\"M226.16,-1270.37C228.68,-1262.16 232.56,-1249.54 235.79,-1239.05\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"239.24,-1239.75 238.83,-1229.17 232.55,-1237.7 239.24,-1239.75\"/>\n</g>\n<!-- 140025091524928 -->\n<g id=\"node26\" class=\"node\">\n<title>140025091524928</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"210,-767 103,-767 103,-726 210,-726 210,-767\"/>\n<text text-anchor=\"middle\" x=\"156.5\" y=\"-755\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n<text text-anchor=\"middle\" x=\"156.5\" y=\"-744\" font-family=\"monospace\" font-size=\"10.00\">step1.fc.weight</text>\n<text text-anchor=\"middle\" x=\"156.5\" y=\"-733\" font-family=\"monospace\" font-size=\"10.00\">(1, 16)</text>\n</g>\n<!-- 140025091526224&#45;&gt;140025091524928 -->\n<g id=\"edge28\" class=\"edge\">\n<title>140025091526224&#45;&gt;140025091524928</title>\n<path fill=\"none\" stroke=\"black\" d=\"M208.81,-1270.48C182.38,-1253.91 129.5,-1214.62 129.5,-1165.5 129.5,-1165.5 129.5,-1165.5 129.5,-888.5 129.5,-849.43 139.62,-805.34 147.47,-777.03\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"150.86,-777.93 150.24,-767.35 144.12,-776 150.86,-777.93\"/>\n</g>\n<!-- 140028155952880 -->\n<g id=\"node18\" class=\"node\">\n<title>140028155952880</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"277,-1361 170,-1361 170,-1331 277,-1331 277,-1361\"/>\n<text text-anchor=\"middle\" x=\"223.5\" y=\"-1349\" font-family=\"monospace\" font-size=\"10.00\">step0.fc.weight</text>\n<text text-anchor=\"middle\" x=\"223.5\" y=\"-1338\" font-family=\"monospace\" font-size=\"10.00\">(1, 16)</text>\n</g>\n<!-- 140028155952880&#45;&gt;140025091526224 -->\n<g id=\"edge17\" class=\"edge\">\n<title>140028155952880&#45;&gt;140025091526224</title>\n<path fill=\"none\" stroke=\"black\" d=\"M223.5,-1330.8C223.5,-1321.7 223.5,-1309.79 223.5,-1299.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"227,-1299.84 223.5,-1289.84 220,-1299.84 227,-1299.84\"/>\n</g>\n<!-- 140025091726784 -->\n<g id=\"node19\" class=\"node\">\n<title>140025091726784</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"281,-415 192,-415 192,-396 281,-396 281,-415\"/>\n<text text-anchor=\"middle\" x=\"236.5\" y=\"-403\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140025091726784&#45;&gt;140027829366544 -->\n<g id=\"edge18\" class=\"edge\">\n<title>140025091726784&#45;&gt;140027829366544</title>\n<path fill=\"none\" stroke=\"black\" d=\"M236.5,-395.75C236.5,-388.8 236.5,-378.85 236.5,-370.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"240,-370.09 236.5,-360.09 233,-370.09 240,-370.09\"/>\n</g>\n<!-- 140025091726688 -->\n<g id=\"node20\" class=\"node\">\n<title>140025091726688</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"284,-470 189,-470 189,-451 284,-451 284,-470\"/>\n<text text-anchor=\"middle\" x=\"236.5\" y=\"-458\" font-family=\"monospace\" font-size=\"10.00\">ViewBackward0</text>\n</g>\n<!-- 140025091726688&#45;&gt;140025091726784 -->\n<g id=\"edge19\" class=\"edge\">\n<title>140025091726688&#45;&gt;140025091726784</title>\n<path fill=\"none\" stroke=\"black\" d=\"M236.5,-450.75C236.5,-443.8 236.5,-433.85 236.5,-425.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"240,-425.09 236.5,-415.09 233,-425.09 240,-425.09\"/>\n</g>\n<!-- 140025091725680 -->\n<g id=\"node21\" class=\"node\">\n<title>140025091725680</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"278,-525 189,-525 189,-506 278,-506 278,-525\"/>\n<text text-anchor=\"middle\" x=\"233.5\" y=\"-513\" font-family=\"monospace\" font-size=\"10.00\">SumBackward1</text>\n</g>\n<!-- 140025091725680&#45;&gt;140025091726688 -->\n<g id=\"edge20\" class=\"edge\">\n<title>140025091725680&#45;&gt;140025091726688</title>\n<path fill=\"none\" stroke=\"black\" d=\"M234,-505.75C234.39,-498.8 234.95,-488.85 235.45,-480.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"238.94,-480.27 236.01,-470.09 231.95,-479.88 238.94,-480.27\"/>\n</g>\n<!-- 140025091726112 -->\n<g id=\"node22\" class=\"node\">\n<title>140025091726112</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"252,-580 91,-580 91,-561 252,-561 252,-580\"/>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-568\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackwardBackward0</text>\n</g>\n<!-- 140025091726112&#45;&gt;140025091725680 -->\n<g id=\"edge21\" class=\"edge\">\n<title>140025091726112&#45;&gt;140025091725680</title>\n<path fill=\"none\" stroke=\"black\" d=\"M181.46,-560.98C190.63,-553.15 204.44,-541.34 215.53,-531.86\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"218.03,-534.33 223.36,-525.17 213.48,-529.01 218.03,-534.33\"/>\n</g>\n<!-- 140025091726880 -->\n<g id=\"node38\" class=\"node\">\n<title>140025091726880</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"169,-525 92,-525 92,-506 169,-506 169,-525\"/>\n<text text-anchor=\"middle\" x=\"130.5\" y=\"-513\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091726112&#45;&gt;140025091726880 -->\n<g id=\"edge43\" class=\"edge\">\n<title>140025091726112&#45;&gt;140025091726880</title>\n<path fill=\"none\" stroke=\"black\" d=\"M164.73,-560.75C158.94,-553.26 150.44,-542.28 143.36,-533.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"146.03,-530.86 137.14,-525.09 140.49,-535.14 146.03,-530.86\"/>\n</g>\n<!-- 140025091726064&#45;&gt;140025091726112 -->\n<g id=\"edge22\" class=\"edge\">\n<title>140025091726064&#45;&gt;140025091726112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M56.04,-1099.82C57.13,-1081.49 59.5,-1037.46 59.5,-1000.5 59.5,-1000.5 59.5,-1000.5 59.5,-745.5 59.5,-682.79 77.36,-665.78 115.5,-616 124.56,-604.18 137.42,-593.66 148.53,-585.83\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"150.63,-588.64 156.97,-580.15 146.72,-582.83 150.63,-588.64\"/>\n</g>\n<!-- 140025091725152&#45;&gt;140025091726112 -->\n<g id=\"edge24\" class=\"edge\">\n<title>140025091725152&#45;&gt;140025091726112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M174.84,-615.75C174.32,-608.8 173.56,-598.85 172.91,-590.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"176.39,-589.8 172.15,-580.09 169.41,-590.32 176.39,-589.8\"/>\n</g>\n<!-- 140025091725824 -->\n<g id=\"node25\" class=\"node\">\n<title>140025091725824</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"210,-690 133,-690 133,-671 210,-671 210,-690\"/>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-678\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091725824&#45;&gt;140025091725152 -->\n<g id=\"edge26\" class=\"edge\">\n<title>140025091725824&#45;&gt;140025091725152</title>\n<path fill=\"none\" stroke=\"black\" d=\"M172.16,-670.75C172.68,-663.8 173.44,-653.85 174.09,-645.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"177.59,-645.32 174.85,-635.09 170.61,-644.8 177.59,-645.32\"/>\n</g>\n<!-- 140025091524928&#45;&gt;140025091725824 -->\n<g id=\"edge27\" class=\"edge\">\n<title>140025091524928&#45;&gt;140025091725824</title>\n<path fill=\"none\" stroke=\"black\" d=\"M161.08,-725.95C163,-717.76 165.22,-708.28 167.12,-700.19\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"170.57,-700.8 169.45,-690.26 163.76,-699.2 170.57,-700.8\"/>\n</g>\n<!-- 140025091726016 -->\n<g id=\"node33\" class=\"node\">\n<title>140025091726016</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"176,-250 87,-250 87,-231 176,-231 176,-250\"/>\n<text text-anchor=\"middle\" x=\"131.5\" y=\"-238\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n</g>\n<!-- 140025091524928&#45;&gt;140025091726016 -->\n<g id=\"edge37\" class=\"edge\">\n<title>140025091524928&#45;&gt;140025091726016</title>\n<path fill=\"none\" stroke=\"black\" d=\"M136.43,-725.86C107.76,-695.59 58.5,-634.43 58.5,-571.5 58.5,-571.5 58.5,-571.5 58.5,-349.5 58.5,-320.04 61.21,-310.54 77.5,-286 85.37,-274.14 97.36,-263.81 108.03,-256.1\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"110.3,-258.79 116.58,-250.26 106.35,-253.02 110.3,-258.79\"/>\n</g>\n<!-- 140025091525600 -->\n<g id=\"node27\" class=\"node\">\n<title>140025091525600</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"247,-833 158,-833 158,-814 247,-814 247,-833\"/>\n<text text-anchor=\"middle\" x=\"202.5\" y=\"-821\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140025091525600&#45;&gt;140025091524928 -->\n<g id=\"edge29\" class=\"edge\">\n<title>140025091525600&#45;&gt;140025091524928</title>\n<path fill=\"none\" stroke=\"black\" d=\"M197.22,-813.9C191.53,-804.62 182.17,-789.35 173.85,-775.79\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"176.8,-773.9 168.59,-767.21 170.83,-777.56 176.8,-773.9\"/>\n</g>\n<!-- 140024973742144 -->\n<g id=\"node28\" class=\"node\">\n<title>140024973742144</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"242,-899 165,-899 165,-880 242,-880 242,-899\"/>\n<text text-anchor=\"middle\" x=\"203.5\" y=\"-887\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973742144&#45;&gt;140025091525600 -->\n<g id=\"edge30\" class=\"edge\">\n<title>140024973742144&#45;&gt;140025091525600</title>\n<path fill=\"none\" stroke=\"black\" d=\"M203.37,-879.87C203.22,-870.66 202.99,-855.79 202.8,-843.77\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"206.3,-843.35 202.64,-833.41 199.3,-843.46 206.3,-843.35\"/>\n</g>\n<!-- 140024973742576 -->\n<g id=\"node29\" class=\"node\">\n<title>140024973742576</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"242,-954 165,-954 165,-935 242,-935 242,-954\"/>\n<text text-anchor=\"middle\" x=\"203.5\" y=\"-942\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973742576&#45;&gt;140024973742144 -->\n<g id=\"edge31\" class=\"edge\">\n<title>140024973742576&#45;&gt;140024973742144</title>\n<path fill=\"none\" stroke=\"black\" d=\"M203.5,-934.75C203.5,-927.8 203.5,-917.85 203.5,-909.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"207,-909.09 203.5,-899.09 200,-909.09 207,-909.09\"/>\n</g>\n<!-- 140024973742480 -->\n<g id=\"node30\" class=\"node\">\n<title>140024973742480</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"246,-1009 163,-1009 163,-990 246,-990 246,-1009\"/>\n<text text-anchor=\"middle\" x=\"204.5\" y=\"-997\" font-family=\"monospace\" font-size=\"10.00\">MmBackward0</text>\n</g>\n<!-- 140024973742480&#45;&gt;140024973742576 -->\n<g id=\"edge32\" class=\"edge\">\n<title>140024973742480&#45;&gt;140024973742576</title>\n<path fill=\"none\" stroke=\"black\" d=\"M204.33,-989.75C204.2,-982.8 204.02,-972.85 203.85,-964.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"207.35,-964.02 203.66,-954.09 200.35,-964.15 207.35,-964.02\"/>\n</g>\n<!-- 140024973742672&#45;&gt;140024973742480 -->\n<g id=\"edge33\" class=\"edge\">\n<title>140024973742672&#45;&gt;140024973742480</title>\n<path fill=\"none\" stroke=\"black\" d=\"M206.17,-1044.75C205.91,-1037.8 205.53,-1027.85 205.2,-1019.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"208.7,-1018.95 204.82,-1009.09 201.7,-1019.21 208.7,-1018.95\"/>\n</g>\n<!-- 140027829365632 -->\n<g id=\"node32\" class=\"node\">\n<title>140027829365632</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"188,-195 111,-195 111,-176 188,-176 188,-195\"/>\n<text text-anchor=\"middle\" x=\"149.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140027829365632&#45;&gt;140027829363616 -->\n<g id=\"edge35\" class=\"edge\">\n<title>140027829365632&#45;&gt;140027829363616</title>\n<path fill=\"none\" stroke=\"black\" d=\"M154.95,-175.75C159.51,-168.42 166.17,-157.73 171.79,-148.7\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"174.84,-150.43 177.15,-140.09 168.9,-146.73 174.84,-150.43\"/>\n</g>\n<!-- 140025091726016&#45;&gt;140027829365632 -->\n<g id=\"edge36\" class=\"edge\">\n<title>140025091726016&#45;&gt;140027829365632</title>\n<path fill=\"none\" stroke=\"black\" d=\"M134.47,-230.75C136.86,-223.72 140.29,-213.62 143.27,-204.84\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"146.68,-205.68 146.58,-195.09 140.05,-203.43 146.68,-205.68\"/>\n</g>\n<!-- 140025091726544 -->\n<g id=\"node34\" class=\"node\">\n<title>140025091726544</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"176,-305 87,-305 87,-286 176,-286 176,-305\"/>\n<text text-anchor=\"middle\" x=\"131.5\" y=\"-293\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140025091726544&#45;&gt;140025091726016 -->\n<g id=\"edge38\" class=\"edge\">\n<title>140025091726544&#45;&gt;140025091726016</title>\n<path fill=\"none\" stroke=\"black\" d=\"M131.5,-285.75C131.5,-278.8 131.5,-268.85 131.5,-260.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"135,-260.09 131.5,-250.09 128,-260.09 135,-260.09\"/>\n</g>\n<!-- 140025091726448 -->\n<g id=\"node35\" class=\"node\">\n<title>140025091726448</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"169,-360 92,-360 92,-341 169,-341 169,-360\"/>\n<text text-anchor=\"middle\" x=\"130.5\" y=\"-348\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091726448&#45;&gt;140025091726544 -->\n<g id=\"edge39\" class=\"edge\">\n<title>140025091726448&#45;&gt;140025091726544</title>\n<path fill=\"none\" stroke=\"black\" d=\"M130.67,-340.75C130.8,-333.8 130.98,-323.85 131.15,-315.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"134.65,-315.15 131.34,-305.09 127.65,-315.02 134.65,-315.15\"/>\n</g>\n<!-- 140025091725584 -->\n<g id=\"node36\" class=\"node\">\n<title>140025091725584</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"168,-415 91,-415 91,-396 168,-396 168,-415\"/>\n<text text-anchor=\"middle\" x=\"129.5\" y=\"-403\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140025091725584&#45;&gt;140025091726448 -->\n<g id=\"edge40\" class=\"edge\">\n<title>140025091725584&#45;&gt;140025091726448</title>\n<path fill=\"none\" stroke=\"black\" d=\"M129.67,-395.75C129.8,-388.8 129.98,-378.85 130.15,-370.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"133.65,-370.15 130.34,-360.09 126.65,-370.02 133.65,-370.15\"/>\n</g>\n<!-- 140025091727024 -->\n<g id=\"node37\" class=\"node\">\n<title>140025091727024</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"171,-470 88,-470 88,-451 171,-451 171,-470\"/>\n<text text-anchor=\"middle\" x=\"129.5\" y=\"-458\" font-family=\"monospace\" font-size=\"10.00\">MmBackward0</text>\n</g>\n<!-- 140025091727024&#45;&gt;140025091725584 -->\n<g id=\"edge41\" class=\"edge\">\n<title>140025091727024&#45;&gt;140025091725584</title>\n<path fill=\"none\" stroke=\"black\" d=\"M129.5,-450.75C129.5,-443.8 129.5,-433.85 129.5,-425.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"133,-425.09 129.5,-415.09 126,-425.09 133,-425.09\"/>\n</g>\n<!-- 140025091726880&#45;&gt;140025091727024 -->\n<g id=\"edge42\" class=\"edge\">\n<title>140025091726880&#45;&gt;140025091727024</title>\n<path fill=\"none\" stroke=\"black\" d=\"M130.33,-505.75C130.2,-498.8 130.02,-488.85 129.85,-480.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"133.35,-480.02 129.66,-470.09 126.35,-480.15 133.35,-480.02\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ───────────────────────────────────────╮</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/tmp/ipykernel_3962266/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">4178930003.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">21</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;module&gt;</span>                                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000; font-style: italic\">[Errno 2] No such file or directory: '/tmp/ipykernel_3962266/4178930003.py'</span>                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/PanXuehai/Miniconda3/envs/torchopt/lib/python3.9/site-packages/torch/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">_tensor.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">487</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">backward</span>           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 484 │   │   │   │   </span>create_graph=create_graph,                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 485 │   │   │   │   </span>inputs=inputs,                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 486 │   │   │   </span>)                                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 487 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>torch.autograd.backward(                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 488 │   │   │   </span><span style=\"color: #00ffff; text-decoration-color: #00ffff\">self</span>, gradient, retain_graph, create_graph, inputs=inputs                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 489 │   │   </span>)                                                                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 490 </span>                                                                                                         <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">╭───────────────────────── locals ──────────────────────────╮</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span> create_graph = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">False</span>                                      <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>     gradient = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                       <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>       inputs = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                       <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span> retain_graph = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                       <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>         self = <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">0.1203</span>, <span style=\"color: #808000; text-decoration-color: #808000\">grad_fn</span>=<span style=\"font-weight: bold\">&lt;</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">MseLossBackward0</span><span style=\"font-weight: bold\">&gt;)</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">╰───────────────────────────────────────────────────────────╯</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/home/PanXuehai/Miniconda3/envs/torchopt/lib/python3.9/site-packages/torch/autograd/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">__init__.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">197</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">backward</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">194 │   # The reason we repeat same the comment below is that</span>                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">195 │   # some Python versions print out the first line of a multi-line function</span>                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">196 │   # calls in the traceback and some print out the last line</span>                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>197 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   </span>Variable._execution_engine.run_backward(  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># Calls into the C++ engine to run the ba</span>                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">198 │   │   </span>tensors, grad_tensors_, retain_graph, create_graph, inputs,                                       <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">199 │   │   </span>allow_unreachable=<span style=\"color: #0000ff; text-decoration-color: #0000ff\">True</span>, accumulate_grad=<span style=\"color: #0000ff; text-decoration-color: #0000ff\">True</span>)  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># Calls into the C++ engine to r</span>                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">200 </span>                                                                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">╭──────────────────────────── locals ────────────────────────────╮</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>   create_graph = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">False</span>                                         <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>   grad_tensors = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                          <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>  grad_tensors_ = <span style=\"font-weight: bold\">(</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>.<span style=\"font-weight: bold\">)</span>,<span style=\"font-weight: bold\">)</span>                                 <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span> grad_variables = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">None</span>                                          <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>         inputs = <span style=\"font-weight: bold\">()</span>                                            <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>   retain_graph = <span style=\"color: #0000ff; text-decoration-color: #0000ff\">False</span>                                         <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>        tensors = <span style=\"font-weight: bold\">(</span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">tensor</span><span style=\"font-weight: bold\">(</span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">0.1203</span>, <span style=\"color: #808000; text-decoration-color: #808000\">grad_fn</span>=<span style=\"font-weight: bold\">&lt;</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">MseLossBackward0</span><span style=\"font-weight: bold\">&gt;)</span>,<span style=\"font-weight: bold\">)</span> <span style=\"color: #808000; text-decoration-color: #808000\">│</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000\">╰────────────────────────────────────────────────────────────────╯</span>                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">RuntimeError: </span>Trying to backward through the graph a second time <span style=\"font-weight: bold\">(</span>or directly access saved tensors after they have \n",
       "already been freed<span style=\"font-weight: bold\">)</span>. Saved intermediate values of the graph are freed when you call <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">.backward</span><span style=\"font-weight: bold\">()</span> or <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">autograd.grad</span><span style=\"font-weight: bold\">()</span>.\n",
       "Specify <span style=\"color: #808000; text-decoration-color: #808000\">retain_graph</span>=<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-style: italic\">True</span> if you need to backward through the graph a second time or if you need to access saved \n",
       "tensors after calling backward.\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[31m╭─\u001b[0m\u001b[31m────────────────────────────────────── \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m ──────────────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/tmp/ipykernel_3962266/\u001b[0m\u001b[1;33m4178930003.py\u001b[0m:\u001b[94m21\u001b[0m in \u001b[92m<module>\u001b[0m                                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[3;31m[Errno 2] No such file or directory: '/tmp/ipykernel_3962266/4178930003.py'\u001b[0m                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/PanXuehai/Miniconda3/envs/torchopt/lib/python3.9/site-packages/torch/\u001b[0m\u001b[1;33m_tensor.py\u001b[0m:\u001b[94m487\u001b[0m in \u001b[92mbackward\u001b[0m           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 484 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mcreate_graph=create_graph,                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 485 \u001b[0m\u001b[2m│   │   │   │   \u001b[0minputs=inputs,                                                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 486 \u001b[0m\u001b[2m│   │   │   \u001b[0m)                                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 487 \u001b[2m│   │   \u001b[0mtorch.autograd.backward(                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 488 \u001b[0m\u001b[2m│   │   │   \u001b[0m\u001b[96mself\u001b[0m, gradient, retain_graph, create_graph, inputs=inputs                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 489 \u001b[0m\u001b[2m│   │   \u001b[0m)                                                                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 490 \u001b[0m                                                                                                         \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m╭─\u001b[0m\u001b[33m──────────────────────── locals ─────────────────────────\u001b[0m\u001b[33m─╮\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m create_graph = \u001b[94mFalse\u001b[0m                                      \u001b[33m│\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m     gradient = \u001b[94mNone\u001b[0m                                       \u001b[33m│\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m       inputs = \u001b[94mNone\u001b[0m                                       \u001b[33m│\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m retain_graph = \u001b[94mNone\u001b[0m                                       \u001b[33m│\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m         self = \u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[94m0.1203\u001b[0m, \u001b[33mgrad_fn\u001b[0m=\u001b[1m<\u001b[0m\u001b[1;95mMseLossBackward0\u001b[0m\u001b[1m>\u001b[0m\u001b[1m)\u001b[0m \u001b[33m│\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m╰───────────────────────────────────────────────────────────╯\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/home/PanXuehai/Miniconda3/envs/torchopt/lib/python3.9/site-packages/torch/autograd/\u001b[0m\u001b[1;33m__init__.py\u001b[0m:\u001b[94m197\u001b[0m in \u001b[92mbackward\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m194 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# The reason we repeat same the comment below is that\u001b[0m                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m195 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# some Python versions print out the first line of a multi-line function\u001b[0m                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m196 \u001b[0m\u001b[2m│   \u001b[0m\u001b[2m# calls in the traceback and some print out the last line\u001b[0m                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m197 \u001b[2m│   \u001b[0mVariable._execution_engine.run_backward(  \u001b[2m# Calls into the C++ engine to run the ba\u001b[0m                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m198 \u001b[0m\u001b[2m│   │   \u001b[0mtensors, grad_tensors_, retain_graph, create_graph, inputs,                                       \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m199 \u001b[0m\u001b[2m│   │   \u001b[0mallow_unreachable=\u001b[94mTrue\u001b[0m, accumulate_grad=\u001b[94mTrue\u001b[0m)  \u001b[2m# Calls into the C++ engine to r\u001b[0m                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m200 \u001b[0m                                                                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m╭─\u001b[0m\u001b[33m─────────────────────────── locals ───────────────────────────\u001b[0m\u001b[33m─╮\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m   create_graph = \u001b[94mFalse\u001b[0m                                         \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m   grad_tensors = \u001b[94mNone\u001b[0m                                          \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m  grad_tensors_ = \u001b[1m(\u001b[0m\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[94m1\u001b[0m.\u001b[1m)\u001b[0m,\u001b[1m)\u001b[0m                                 \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m grad_variables = \u001b[94mNone\u001b[0m                                          \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m         inputs = \u001b[1m(\u001b[0m\u001b[1m)\u001b[0m                                            \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m   retain_graph = \u001b[94mFalse\u001b[0m                                         \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m│\u001b[0m        tensors = \u001b[1m(\u001b[0m\u001b[1;35mtensor\u001b[0m\u001b[1m(\u001b[0m\u001b[94m0.1203\u001b[0m, \u001b[33mgrad_fn\u001b[0m=\u001b[1m<\u001b[0m\u001b[1;95mMseLossBackward0\u001b[0m\u001b[1m>\u001b[0m\u001b[1m)\u001b[0m,\u001b[1m)\u001b[0m \u001b[33m│\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[33m╰────────────────────────────────────────────────────────────────╯\u001b[0m                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
       "\u001b[1;91mRuntimeError: \u001b[0mTrying to backward through the graph a second time \u001b[1m(\u001b[0mor directly access saved tensors after they have \n",
       "already been freed\u001b[1m)\u001b[0m. Saved intermediate values of the graph are freed when you call \u001b[1;35m.backward\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m or \u001b[1;35mautograd.grad\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m.\n",
       "Specify \u001b[33mretain_graph\u001b[0m=\u001b[3;92mTrue\u001b[0m if you need to backward through the graph a second time or if you need to access saved \n",
       "tensors after calling backward.\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Inner update with attached computation graph\n",
    "inner_loss = loss_fn(net(x), y)\n",
    "loss = inner_loss * meta_parameter\n",
    "optim.step(loss)\n",
    "\n",
    "# Outer forward process\n",
    "outer_loss = loss_fn(net(x), y)\n",
    "display(\n",
    "    torchopt.visual.make_dot(\n",
    "        outer_loss,\n",
    "        params=(\n",
    "            init_net_state,\n",
    "            one_step_net_state,\n",
    "            {'meta_parameter': meta_parameter, 'outer_loss': outer_loss},\n",
    "        ),\n",
    "    )\n",
    ")\n",
    "\n",
    "# Outer update\n",
    "meta_optim.zero_grad()\n",
    "outer_loss.backward()\n",
    "meta_optim.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the graph we can see, directly conducting the second bi-level process links the graph of first and second bi-level process together. We should manually stop gradient with `torchopt.stop_gradient`. `torchopt.stop_gradient` will detach the node of gradient graph and make it become a leaf node. It allows the input of network, optimizer, or state dictionary and the gradient operation happens in an in-place manner.\n",
    "\n",
    "Let's use `recover_state_dict` to come back to one-step updated states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset to previous one-step updated states\n",
    "torchopt.recover_state_dict(net, one_step_net_state)\n",
    "torchopt.recover_state_dict(optim, one_step_optim_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, Let's conduct the stop-gradient operation before the second meta-optimization step. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "meta_parameter.grad = tensor(-0.0635)\n",
      "meta_parameter = Parameter containing:\n",
      "tensor(1.1940, requires_grad=True)\n",
      "<graphviz.graphs.Digraph object at 0x7f5a19ced640>\n"
     ]
    },
    {
     "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.42.3 (20191010.1750)\n -->\n<!-- Title: %3 Pages: 1 -->\n<svg width=\"437pt\" height=\"830pt\"\n viewBox=\"0.00 0.00 437.00 830.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 826)\">\n<title>%3</title>\n<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-826 433,-826 433,4 -4,4\"/>\n<!-- 140024973754912 -->\n<g id=\"node1\" class=\"node\">\n<title>140024973754912</title>\n<polygon fill=\"#caff70\" stroke=\"black\" points=\"210,-30 133,-30 133,0 210,0 210,-30\"/>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-18\" font-family=\"monospace\" font-size=\"10.00\">outer_loss</text>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-7\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140024956770528 -->\n<g id=\"node2\" class=\"node\">\n<title>140024956770528</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"228,-85 115,-85 115,-66 228,-66 228,-85\"/>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-73\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackward0</text>\n</g>\n<!-- 140024956770528&#45;&gt;140024973754912 -->\n<g id=\"edge26\" class=\"edge\">\n<title>140024956770528&#45;&gt;140024973754912</title>\n<path fill=\"none\" stroke=\"black\" d=\"M171.5,-65.87C171.5,-59.11 171.5,-49.35 171.5,-40.26\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"175,-40.11 171.5,-30.11 168,-40.11 175,-40.11\"/>\n</g>\n<!-- 140024956772112 -->\n<g id=\"node3\" class=\"node\">\n<title>140024956772112</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"222,-140 121,-140 121,-121 222,-121 222,-140\"/>\n<text text-anchor=\"middle\" x=\"171.5\" y=\"-128\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140024956772112&#45;&gt;140024956770528 -->\n<g id=\"edge1\" class=\"edge\">\n<title>140024956772112&#45;&gt;140024956770528</title>\n<path fill=\"none\" stroke=\"black\" d=\"M171.5,-120.75C171.5,-113.8 171.5,-103.85 171.5,-95.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"175,-95.09 171.5,-85.09 168,-95.09 175,-95.09\"/>\n</g>\n<!-- 140024956770720 -->\n<g id=\"node4\" class=\"node\">\n<title>140024956770720</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"167,-360 78,-360 78,-341 167,-341 167,-360\"/>\n<text text-anchor=\"middle\" x=\"122.5\" y=\"-348\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n</g>\n<!-- 140024956770720&#45;&gt;140024956772112 -->\n<g id=\"edge2\" class=\"edge\">\n<title>140024956770720&#45;&gt;140024956772112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M124.32,-340.82C129.65,-315.41 145.8,-239.06 160.5,-176 162.51,-167.37 164.86,-157.82 166.86,-149.8\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"170.28,-150.57 169.32,-140.01 163.49,-148.86 170.28,-150.57\"/>\n</g>\n<!-- 140024962101312 -->\n<g id=\"node5\" class=\"node\">\n<title>140024962101312</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"125,-690 24,-690 24,-671 125,-671 125,-690\"/>\n<text text-anchor=\"middle\" x=\"74.5\" y=\"-678\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140024962101312&#45;&gt;140024956770720 -->\n<g id=\"edge3\" class=\"edge\">\n<title>140024962101312&#45;&gt;140024956770720</title>\n<path fill=\"none\" stroke=\"black\" d=\"M70.61,-670.74C62.91,-652.6 46.5,-609.46 46.5,-571.5 46.5,-571.5 46.5,-571.5 46.5,-459.5 46.5,-429.63 51.45,-420.52 68.5,-396 76.73,-384.16 88.95,-373.72 99.68,-365.94\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"101.97,-368.61 108.22,-360.06 98,-362.85 101.97,-368.61\"/>\n</g>\n<!-- 140024973745552 -->\n<g id=\"node14\" class=\"node\">\n<title>140024973745552</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"241,-635 140,-635 140,-616 241,-616 241,-635\"/>\n<text text-anchor=\"middle\" x=\"190.5\" y=\"-623\" font-family=\"monospace\" font-size=\"10.00\">AddmmBackward0</text>\n</g>\n<!-- 140024962101312&#45;&gt;140024973745552 -->\n<g id=\"edge13\" class=\"edge\">\n<title>140024962101312&#45;&gt;140024973745552</title>\n<path fill=\"none\" stroke=\"black\" d=\"M93.14,-670.98C111.8,-662.46 140.75,-649.23 162.24,-639.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"163.88,-642.51 171.52,-635.17 160.97,-636.14 163.88,-642.51\"/>\n</g>\n<!-- 140025091547520 -->\n<g id=\"node6\" class=\"node\">\n<title>140025091547520</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"149,-756 0,-756 0,-726 149,-726 149,-756\"/>\n<text text-anchor=\"middle\" x=\"74.5\" y=\"-744\" font-family=\"monospace\" font-size=\"10.00\">step1.detached.fc.bias</text>\n<text text-anchor=\"middle\" x=\"74.5\" y=\"-733\" font-family=\"monospace\" font-size=\"10.00\">(1)</text>\n</g>\n<!-- 140025091547520&#45;&gt;140024962101312 -->\n<g id=\"edge4\" class=\"edge\">\n<title>140025091547520&#45;&gt;140024962101312</title>\n<path fill=\"none\" stroke=\"black\" d=\"M74.5,-725.84C74.5,-718.21 74.5,-708.7 74.5,-700.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"78,-700.27 74.5,-690.27 71,-700.27 78,-700.27\"/>\n</g>\n<!-- 140024971586864 -->\n<g id=\"node7\" class=\"node\">\n<title>140024971586864</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"167,-415 78,-415 78,-396 167,-396 167,-415\"/>\n<text text-anchor=\"middle\" x=\"122.5\" y=\"-403\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140024971586864&#45;&gt;140024956770720 -->\n<g id=\"edge5\" class=\"edge\">\n<title>140024971586864&#45;&gt;140024956770720</title>\n<path fill=\"none\" stroke=\"black\" d=\"M122.5,-395.75C122.5,-388.8 122.5,-378.85 122.5,-370.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"126,-370.09 122.5,-360.09 119,-370.09 126,-370.09\"/>\n</g>\n<!-- 140024973742528 -->\n<g id=\"node8\" class=\"node\">\n<title>140024973742528</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"170,-470 75,-470 75,-451 170,-451 170,-470\"/>\n<text text-anchor=\"middle\" x=\"122.5\" y=\"-458\" font-family=\"monospace\" font-size=\"10.00\">ViewBackward0</text>\n</g>\n<!-- 140024973742528&#45;&gt;140024971586864 -->\n<g id=\"edge6\" class=\"edge\">\n<title>140024973742528&#45;&gt;140024971586864</title>\n<path fill=\"none\" stroke=\"black\" d=\"M122.5,-450.75C122.5,-443.8 122.5,-433.85 122.5,-425.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"126,-425.09 122.5,-415.09 119,-425.09 126,-425.09\"/>\n</g>\n<!-- 140024973743968 -->\n<g id=\"node9\" class=\"node\">\n<title>140024973743968</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"168,-525 79,-525 79,-506 168,-506 168,-525\"/>\n<text text-anchor=\"middle\" x=\"123.5\" y=\"-513\" font-family=\"monospace\" font-size=\"10.00\">SumBackward1</text>\n</g>\n<!-- 140024973743968&#45;&gt;140024973742528 -->\n<g id=\"edge7\" class=\"edge\">\n<title>140024973743968&#45;&gt;140024973742528</title>\n<path fill=\"none\" stroke=\"black\" d=\"M123.33,-505.75C123.2,-498.8 123.02,-488.85 122.85,-480.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"126.35,-480.02 122.66,-470.09 119.35,-480.15 126.35,-480.02\"/>\n</g>\n<!-- 140024973742768 -->\n<g id=\"node10\" class=\"node\">\n<title>140024973742768</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"271,-580 110,-580 110,-561 271,-561 271,-580\"/>\n<text text-anchor=\"middle\" x=\"190.5\" y=\"-568\" font-family=\"monospace\" font-size=\"10.00\">MseLossBackwardBackward0</text>\n</g>\n<!-- 140024973742768&#45;&gt;140024973743968 -->\n<g id=\"edge8\" class=\"edge\">\n<title>140024973742768&#45;&gt;140024973743968</title>\n<path fill=\"none\" stroke=\"black\" d=\"M179.74,-560.98C169.73,-553.07 154.61,-541.11 142.57,-531.58\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"144.47,-528.63 134.46,-525.17 140.13,-534.12 144.47,-528.63\"/>\n</g>\n<!-- 140024973744400 -->\n<g id=\"node24\" class=\"node\">\n<title>140024973744400</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"265,-525 188,-525 188,-506 265,-506 265,-525\"/>\n<text text-anchor=\"middle\" x=\"226.5\" y=\"-513\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973742768&#45;&gt;140024973744400 -->\n<g id=\"edge25\" class=\"edge\">\n<title>140024973742768&#45;&gt;140024973744400</title>\n<path fill=\"none\" stroke=\"black\" d=\"M196.44,-560.75C201.48,-553.34 208.84,-542.5 215.01,-533.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"217.94,-535.33 220.67,-525.09 212.15,-531.39 217.94,-535.33\"/>\n</g>\n<!-- 140024973744688 -->\n<g id=\"node11\" class=\"node\">\n<title>140024973744688</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"420,-635 331,-635 331,-616 420,-616 420,-635\"/>\n<text text-anchor=\"middle\" x=\"375.5\" y=\"-623\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140024973744688&#45;&gt;140024973742768 -->\n<g id=\"edge9\" class=\"edge\">\n<title>140024973744688&#45;&gt;140024973742768</title>\n<path fill=\"none\" stroke=\"black\" d=\"M345.78,-615.98C314.45,-607.01 264.93,-592.82 230.26,-582.89\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"231.03,-579.47 220.46,-580.08 229.11,-586.2 231.03,-579.47\"/>\n</g>\n<!-- 140024973745264 -->\n<g id=\"node12\" class=\"node\">\n<title>140024973745264</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"429,-690 328,-690 328,-671 429,-671 429,-690\"/>\n<text text-anchor=\"middle\" x=\"378.5\" y=\"-678\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140024973745264&#45;&gt;140024973744688 -->\n<g id=\"edge10\" class=\"edge\">\n<title>140024973745264&#45;&gt;140024973744688</title>\n<path fill=\"none\" stroke=\"black\" d=\"M378,-670.75C377.61,-663.8 377.05,-653.85 376.55,-645.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"380.05,-644.88 375.99,-635.09 373.06,-645.27 380.05,-644.88\"/>\n</g>\n<!-- 140025091549440 -->\n<g id=\"node13\" class=\"node\">\n<title>140025091549440</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"429,-756 328,-756 328,-726 429,-726 429,-756\"/>\n<text text-anchor=\"middle\" x=\"378.5\" y=\"-744\" font-family=\"monospace\" font-size=\"10.00\">meta_parameter</text>\n<text text-anchor=\"middle\" x=\"378.5\" y=\"-733\" font-family=\"monospace\" font-size=\"10.00\">()</text>\n</g>\n<!-- 140025091549440&#45;&gt;140024973745264 -->\n<g id=\"edge11\" class=\"edge\">\n<title>140025091549440&#45;&gt;140024973745264</title>\n<path fill=\"none\" stroke=\"black\" d=\"M378.5,-725.84C378.5,-718.21 378.5,-708.7 378.5,-700.45\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"382,-700.27 378.5,-690.27 375,-700.27 382,-700.27\"/>\n</g>\n<!-- 140024973745552&#45;&gt;140024973742768 -->\n<g id=\"edge12\" class=\"edge\">\n<title>140024973745552&#45;&gt;140024973742768</title>\n<path fill=\"none\" stroke=\"black\" d=\"M190.5,-615.75C190.5,-608.8 190.5,-598.85 190.5,-590.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"194,-590.09 190.5,-580.09 187,-590.09 194,-590.09\"/>\n</g>\n<!-- 140024973745168 -->\n<g id=\"node15\" class=\"node\">\n<title>140024973745168</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"250,-690 173,-690 173,-671 250,-671 250,-690\"/>\n<text text-anchor=\"middle\" x=\"211.5\" y=\"-678\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973745168&#45;&gt;140024973745552 -->\n<g id=\"edge14\" class=\"edge\">\n<title>140024973745168&#45;&gt;140024973745552</title>\n<path fill=\"none\" stroke=\"black\" d=\"M208.03,-670.75C205.22,-663.65 201.16,-653.4 197.65,-644.56\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"200.84,-643.1 193.9,-635.09 194.33,-645.68 200.84,-643.1\"/>\n</g>\n<!-- 140024973744256 -->\n<g id=\"node16\" class=\"node\">\n<title>140024973744256</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"297,-750.5 196,-750.5 196,-731.5 297,-731.5 297,-750.5\"/>\n<text text-anchor=\"middle\" x=\"246.5\" y=\"-738.5\" font-family=\"monospace\" font-size=\"10.00\">AccumulateGrad</text>\n</g>\n<!-- 140024973744256&#45;&gt;140024973745168 -->\n<g id=\"edge15\" class=\"edge\">\n<title>140024973744256&#45;&gt;140024973745168</title>\n<path fill=\"none\" stroke=\"black\" d=\"M241.34,-731.37C236.27,-722.9 228.4,-709.74 222.01,-699.07\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"224.82,-696.95 216.69,-690.17 218.82,-700.55 224.82,-696.95\"/>\n</g>\n<!-- 140024973745984 -->\n<g id=\"node19\" class=\"node\">\n<title>140024973745984</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"271,-250 182,-250 182,-231 271,-231 271,-250\"/>\n<text text-anchor=\"middle\" x=\"226.5\" y=\"-238\" font-family=\"monospace\" font-size=\"10.00\">AddBackward0</text>\n</g>\n<!-- 140024973744256&#45;&gt;140024973745984 -->\n<g id=\"edge19\" class=\"edge\">\n<title>140024973744256&#45;&gt;140024973745984</title>\n<path fill=\"none\" stroke=\"black\" d=\"M253.87,-731.24C268.43,-713.01 299.5,-669.12 299.5,-626.5 299.5,-626.5 299.5,-626.5 299.5,-349.5 299.5,-319.91 295.93,-310.61 279.5,-286 271.58,-274.13 259.6,-263.7 249.04,-255.93\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"250.82,-252.9 240.61,-250.05 246.81,-258.64 250.82,-252.9\"/>\n</g>\n<!-- 140027828983424 -->\n<g id=\"node17\" class=\"node\">\n<title>140027828983424</title>\n<polygon fill=\"lightblue\" stroke=\"black\" points=\"327,-822 166,-822 166,-792 327,-792 327,-822\"/>\n<text text-anchor=\"middle\" x=\"246.5\" y=\"-810\" font-family=\"monospace\" font-size=\"10.00\">step1.detached.fc.weight</text>\n<text text-anchor=\"middle\" x=\"246.5\" y=\"-799\" font-family=\"monospace\" font-size=\"10.00\">(1, 16)</text>\n</g>\n<!-- 140027828983424&#45;&gt;140024973744256 -->\n<g id=\"edge16\" class=\"edge\">\n<title>140027828983424&#45;&gt;140024973744256</title>\n<path fill=\"none\" stroke=\"black\" d=\"M246.5,-791.8C246.5,-782.7 246.5,-770.79 246.5,-760.9\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"250,-760.84 246.5,-750.84 243,-760.84 250,-760.84\"/>\n</g>\n<!-- 140024956771632 -->\n<g id=\"node18\" class=\"node\">\n<title>140024956771632</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"247,-195 170,-195 170,-176 247,-176 247,-195\"/>\n<text text-anchor=\"middle\" x=\"208.5\" y=\"-183\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024956771632&#45;&gt;140024956772112 -->\n<g id=\"edge17\" class=\"edge\">\n<title>140024956771632&#45;&gt;140024956772112</title>\n<path fill=\"none\" stroke=\"black\" d=\"M202.39,-175.75C197.22,-168.34 189.65,-157.5 183.31,-148.41\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"186.09,-146.29 177.5,-140.09 180.35,-150.29 186.09,-146.29\"/>\n</g>\n<!-- 140024973745984&#45;&gt;140024956771632 -->\n<g id=\"edge18\" class=\"edge\">\n<title>140024973745984&#45;&gt;140024956771632</title>\n<path fill=\"none\" stroke=\"black\" d=\"M223.53,-230.75C221.14,-223.72 217.71,-213.62 214.73,-204.84\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"217.95,-203.43 211.42,-195.09 211.32,-205.68 217.95,-203.43\"/>\n</g>\n<!-- 140024973743728 -->\n<g id=\"node20\" class=\"node\">\n<title>140024973743728</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"271,-305 182,-305 182,-286 271,-286 271,-305\"/>\n<text text-anchor=\"middle\" x=\"226.5\" y=\"-293\" font-family=\"monospace\" font-size=\"10.00\">MulBackward0</text>\n</g>\n<!-- 140024973743728&#45;&gt;140024973745984 -->\n<g id=\"edge20\" class=\"edge\">\n<title>140024973743728&#45;&gt;140024973745984</title>\n<path fill=\"none\" stroke=\"black\" d=\"M226.5,-285.75C226.5,-278.8 226.5,-268.85 226.5,-260.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"230,-260.09 226.5,-250.09 223,-260.09 230,-260.09\"/>\n</g>\n<!-- 140024973743344 -->\n<g id=\"node21\" class=\"node\">\n<title>140024973743344</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"266,-360 189,-360 189,-341 266,-341 266,-360\"/>\n<text text-anchor=\"middle\" x=\"227.5\" y=\"-348\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973743344&#45;&gt;140024973743728 -->\n<g id=\"edge21\" class=\"edge\">\n<title>140024973743344&#45;&gt;140024973743728</title>\n<path fill=\"none\" stroke=\"black\" d=\"M227.33,-340.75C227.2,-333.8 227.02,-323.85 226.85,-315.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"230.35,-315.02 226.66,-305.09 223.35,-315.15 230.35,-315.02\"/>\n</g>\n<!-- 140024973745312 -->\n<g id=\"node22\" class=\"node\">\n<title>140024973745312</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"267,-415 190,-415 190,-396 267,-396 267,-415\"/>\n<text text-anchor=\"middle\" x=\"228.5\" y=\"-403\" font-family=\"monospace\" font-size=\"10.00\">TBackward0</text>\n</g>\n<!-- 140024973745312&#45;&gt;140024973743344 -->\n<g id=\"edge22\" class=\"edge\">\n<title>140024973745312&#45;&gt;140024973743344</title>\n<path fill=\"none\" stroke=\"black\" d=\"M228.33,-395.75C228.2,-388.8 228.02,-378.85 227.85,-370.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"231.35,-370.02 227.66,-360.09 224.35,-370.15 231.35,-370.02\"/>\n</g>\n<!-- 140024973743200 -->\n<g id=\"node23\" class=\"node\">\n<title>140024973743200</title>\n<polygon fill=\"lightgrey\" stroke=\"black\" points=\"271,-470 188,-470 188,-451 271,-451 271,-470\"/>\n<text text-anchor=\"middle\" x=\"229.5\" y=\"-458\" font-family=\"monospace\" font-size=\"10.00\">MmBackward0</text>\n</g>\n<!-- 140024973743200&#45;&gt;140024973745312 -->\n<g id=\"edge23\" class=\"edge\">\n<title>140024973743200&#45;&gt;140024973745312</title>\n<path fill=\"none\" stroke=\"black\" d=\"M229.33,-450.75C229.2,-443.8 229.02,-433.85 228.85,-425.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"232.35,-425.02 228.66,-415.09 225.35,-425.15 232.35,-425.02\"/>\n</g>\n<!-- 140024973744400&#45;&gt;140024973743200 -->\n<g id=\"edge24\" class=\"edge\">\n<title>140024973744400&#45;&gt;140024973743200</title>\n<path fill=\"none\" stroke=\"black\" d=\"M227,-505.75C227.39,-498.8 227.95,-488.85 228.45,-480.13\"/>\n<polygon fill=\"black\" stroke=\"black\" points=\"231.94,-480.27 229.01,-470.09 224.95,-479.88 231.94,-480.27\"/>\n</g>\n</g>\n</svg>\n"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Stop gradient and make them become the leaf node\n",
    "torchopt.stop_gradient(net)\n",
    "torchopt.stop_gradient(optim)\n",
    "one_step_net_state_detached = torchopt.extract_state_dict(\n",
    "    net, enable_visual=True, visual_prefix='step1.detached.'\n",
    ")\n",
    "\n",
    "# Inner update\n",
    "inner_loss = loss_fn(net(x), y)\n",
    "loss = inner_loss * meta_parameter\n",
    "optim.step(loss)\n",
    "\n",
    "# Outer update\n",
    "outer_loss = loss_fn(net(x), y)\n",
    "meta_optim.zero_grad()\n",
    "outer_loss.backward()\n",
    "print(f'meta_parameter.grad = {meta_parameter.grad!r}')\n",
    "meta_optim.step()\n",
    "print(f'meta_parameter = {meta_parameter!r}')\n",
    "\n",
    "display(\n",
    "    torchopt.visual.make_dot(\n",
    "        outer_loss,\n",
    "        params=(\n",
    "            one_step_net_state_detached,\n",
    "            {'meta_parameter': meta_parameter, 'outer_loss': outer_loss},\n",
    "        ),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The gradient graph is the same with the first meta-optimization's gradient graph and we successfully conduct the second bi-level process."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "2a8cc1ff2cbc47027bf9993941710d9ab9175f14080903d9c7c432ee63d681da"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}