{
    "_model__stem":{
       "inputs":[
          "_inputs__1"
       ],
       "outputs":[
          "pool2__1"
       ],
       "nodes":{
          "_model__stem__conv1":{
             "operation":{
                "type":"__torch__.torch.nn.modules.conv.Conv2d",
                "parameters":{
                   "out_channels":32,
                   "in_channels":1,
                   "kernel_size":5
                }
             }
          },
          "_model__stem__pool1":{
             "operation":{
                "type":"__torch__.torch.nn.modules.pooling.MaxPool2d",
                "parameters":{
                   "kernel_size":2
                }
             }
          },
          "_model__stem__conv2":{
             "operation":{
                "type":"__torch__.torch.nn.modules.conv.Conv2d",
                "parameters":{
                   "out_channels":64,
                   "in_channels":32,
                   "kernel_size":5
                }
             }
          },
          "_model__stem__pool2":{
             "operation":{
                "type":"__torch__.torch.nn.modules.pooling.MaxPool2d",
                "parameters":{
                   "kernel_size":2
                }
             }
          }
       },
       "edges":[
          {
             "head":[
                "_inputs",
                0
             ],
             "tail":[
                "_model__stem__conv1",
                0
             ]
          },
          {
             "head":[
                "_model__stem__conv1",
                null
             ],
             "tail":[
                "_model__stem__pool1",
                0
             ]
          },
          {
             "head":[
                "_model__stem__pool1",
                null
             ],
             "tail":[
                "_model__stem__conv2",
                0
             ]
          },
          {
             "head":[
                "_model__stem__conv2",
                null
             ],
             "tail":[
                "_model__stem__pool2",
                0
             ]
          },
          {
             "head":[
                "_model__stem__pool2",
                null
             ],
             "tail":[
                "_outputs",
                null
             ]
          }
       ]
    },
    "_model":{
       "inputs":[
          "image__1"
       ],
       "outputs":[
          "softmax__1"
       ],
       "nodes":{
          "_model__Constant2":{
             "operation":{
                "type":"prim::Constant",
                "parameters":{
                   
                }
             }
          },
          "_model__Constant3":{
             "operation":{
                "type":"prim::Constant",
                "parameters":{
                   "value":3
                }
             }
          },
          "_model__Constant4":{
             "operation":{
                "type":"prim::Constant",
                "parameters":{
                   "value":-1
                }
             }
          },
          "_model__Constant5":{
             "operation":{
                "type":"prim::Constant",
                "parameters":{
                   "value":0
                }
             }
          },
          "_model__stem":{
             "operation":{
                "type":"_cell",
                "parameters":{
                   
                },
                "cell_name":"_model__stem"
             }
          },
          "_model__Size6":{
             "operation":{
                "type":"aten::size",
                "parameters":{
                   
                }
             }
          },
          "_model__ListConstruct7":{
             "operation":{
                "type":"prim::ListConstruct",
                "parameters":{
                   
                }
             }
          },
          "_model__View8":{
             "operation":{
                "type":"aten::view",
                "parameters":{
                   
                }
             }
          },
          "_model__fc1":{
             "operation":{
                "type":"__torch__.torch.nn.modules.linear.Linear",
                "parameters":{
                   "in_features":1024,
                   "out_features":256
                }
             }
          },
          "_model__fc2":{
             "operation":{
                "type":"__torch__.torch.nn.modules.linear.Linear",
                "parameters":{
                   "in_features":256,
                   "out_features":10
                }
             }
          },
          "_model__softmax9":{
             "operation":{
                "type":"Function.softmax",
                "parameters":{
                   
                }
             }
          }
       },
       "edges":[
          {
             "head":[
                "_inputs",
                0
             ],
             "tail":[
                "_model__stem",
                0
             ]
          },
          {
             "head":[
                "_model__stem",
                null
             ],
             "tail":[
                "_model__Size6",
                0
             ]
          },
          {
             "head":[
                "_model__Constant5",
                null
             ],
             "tail":[
                "_model__Size6",
                1
             ]
          },
          {
             "head":[
                "_model__Size6",
                null
             ],
             "tail":[
                "_model__ListConstruct7",
                0
             ]
          },
          {
             "head":[
                "_model__Constant4",
                null
             ],
             "tail":[
                "_model__ListConstruct7",
                1
             ]
          },
          {
             "head":[
                "_model__stem",
                null
             ],
             "tail":[
                "_model__View8",
                0
             ]
          },
          {
             "head":[
                "_model__ListConstruct7",
                null
             ],
             "tail":[
                "_model__View8",
                1
             ]
          },
          {
             "head":[
                "_model__View8",
                null
             ],
             "tail":[
                "_model__fc1",
                0
             ]
          },
          {
             "head":[
                "_model__fc1",
                null
             ],
             "tail":[
                "_model__fc2",
                0
             ]
          },
          {
             "head":[
                "_model__fc2",
                null
             ],
             "tail":[
                "_model__softmax9",
                0
             ]
          },
          {
             "head":[
                "_model__Constant4",
                null
             ],
             "tail":[
                "_model__softmax9",
                1
             ]
          },
          {
             "head":[
                "_model__Constant3",
                null
             ],
             "tail":[
                "_model__softmax9",
                2
             ]
          },
          {
             "head":[
                "_model__Constant2",
                null
             ],
             "tail":[
                "_model__softmax9",
                3
             ]
          },
          {
             "head":[
                "_model__softmax9",
                null
             ],
             "tail":[
                "_outputs",
                null
             ]
          }
       ]
    },
    "_evaluator": {
        "module": "nni.retiarii.trainer.PyTorchImageClassificationTrainer",
        "kwargs": {
            "dataset_cls": "MNIST",
            "dataset_kwargs": {
                "root": "data/mnist",
                "download": true
            },
            "dataloader_kwargs": {
                "batch_size": 32
            },
            "optimizer_cls" : "SGD",
            "optimizer_kwargs": {
                "lr": 1e-3
            },
            "trainer_kwargs": {
                "max_epochs": 1
            }
        }
    }
 }