notebooks/models_postprocessing/jtree/3_train_encoder_predictor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load functions, set variables, create new functions for dataset and model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "import sys\n",
    "sys.path.append('../../../../icml18-jtnn')\n",
    "sys.path.append('../../../../icml18-jtnn/jtnn')\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import math, random, sys\n",
    "from optparse import OptionParser\n",
    "from collections import deque\n",
    "\n",
    "from jtnn import *\n",
    "import rdkit\n",
    "\n",
    "from jtnn_enc import JTNNEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "VOCAB_PATH = '../../../../icml18-jtnn/data/zinc/vocab.txt'\n",
    "MODEL_PATH = '../../../../icml18-jtnn/molvae'\n",
    "DATASET_PATH = '../../../data/3_final_data/split_data'\n",
    "\n",
    "RAW_PATH = '../../../data/raw/baselines/jtree'\n",
    "\n",
    "SMILES_COLUMN = 'smiles'\n",
    "VALUE_COLUMN = 'logP'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = [x.strip(\"\\r\\n \") for x in open(VOCAB_PATH)] \n",
    "vocab = Vocab(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = int(200)\n",
    "hidden_size = int(450)\n",
    "latent_size = int(56)\n",
    "depth = int(3)\n",
    "beta = float(0)\n",
    "lr = float(1e-5)\n",
    "stereo = True if int(1) == 1 else False\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_batch_nodeID(mol_batch, vocab):\n",
    "    tot = 0\n",
    "    for mol_tree in mol_batch:\n",
    "        for node in mol_tree.nodes:\n",
    "            node.idx = tot\n",
    "            node.wid = vocab.get_index(node.smiles)\n",
    "            tot += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class JTPredict(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab, hidden_size, latent_size, depth, stereo=True):\n",
    "        super(JTPredict, self).__init__()\n",
    "        self.vocab = vocab\n",
    "        self.hidden_size = hidden_size\n",
    "        self.latent_size = latent_size\n",
    "        self.depth = depth\n",
    "\n",
    "        self.embedding = nn.Embedding(vocab.size(), hidden_size)\n",
    "        self.jtnn = JTNNEncoder(vocab, hidden_size, self.embedding)\n",
    "        self.mpn = MPN(hidden_size, depth)\n",
    "        \n",
    "        self.output_size = 1\n",
    "\n",
    "        self.T_mean = nn.Linear(hidden_size, latent_size / 2)\n",
    "        self.T_var = nn.Linear(hidden_size, latent_size / 2)\n",
    "        self.G_mean = nn.Linear(hidden_size, latent_size / 2)\n",
    "        self.G_var = nn.Linear(hidden_size, latent_size / 2)\n",
    "        \n",
    "        self.use_stereo = stereo\n",
    "        if stereo:\n",
    "            self.stereo_loss = nn.CrossEntropyLoss(size_average=False)\n",
    "    \n",
    "    def encode(self, mol_batch):\n",
    "        set_batch_nodeID(mol_batch, self.vocab)\n",
    "        root_batch = [mol_tree.nodes[0] for mol_tree in mol_batch]\n",
    "        tree_mess,tree_vec = self.jtnn(root_batch)\n",
    "\n",
    "        smiles_batch = [mol_tree.smiles for mol_tree in mol_batch]\n",
    "        mol_vec = self.mpn(mol2graph(smiles_batch))\n",
    "        return tree_mess, tree_vec, mol_vec\n",
    "\n",
    "    def encode_latent_mean(self, smiles_list):\n",
    "        print(smiles_list)\n",
    "        mol_batch = [MolTree(s) for s in smiles_list]\n",
    "#         print(mol_batch)\n",
    "        for mol_tree in mol_batch:\n",
    "            mol_tree.recover()\n",
    "\n",
    "        _, tree_vec, mol_vec = self.encode(mol_batch)\n",
    "        tree_mean = self.T_mean(tree_vec)\n",
    "        mol_mean = self.G_mean(mol_vec)\n",
    "        return torch.cat([tree_mean,mol_mean], dim=1)\n",
    "    \n",
    "    def create_ffn(self, ffn_num_layers = 3, ffn_hidden_size = 50):\n",
    "        \"\"\"\n",
    "        Creates the feed-forward layers for the model.\n",
    "        :param args: A :class:`~chemprop.args.TrainArgs` object containing model arguments.\n",
    "        \"\"\"\n",
    "        dropout = nn.Dropout(0.5)\n",
    "        activation = nn.ReLU()\n",
    "        \n",
    "        first_linear_dim = self.latent_size\n",
    "\n",
    "        # Create FFN layers\n",
    "        if ffn_num_layers == 1:\n",
    "            ffn = [\n",
    "                dropout,\n",
    "                nn.Linear(first_linear_dim, self.output_size)\n",
    "            ]\n",
    "        else:\n",
    "            ffn = [\n",
    "                dropout,\n",
    "                nn.Linear(first_linear_dim, ffn_hidden_size)\n",
    "            ]\n",
    "            for _ in range(ffn_num_layers - 2):\n",
    "                ffn.extend([\n",
    "                    activation,\n",
    "                    dropout,\n",
    "                    nn.Linear(ffn_hidden_size, ffn_hidden_size),\n",
    "                ])\n",
    "            ffn.extend([\n",
    "                activation,\n",
    "                dropout,\n",
    "                nn.Linear(ffn_hidden_size, self.output_size),\n",
    "            ])\n",
    "\n",
    "        # Create FFN model\n",
    "        self.ffn = nn.Sequential(*ffn)\n",
    "\n",
    "    def forward(self, mol_batch, beta=0):\n",
    "        batch_size = len(mol_batch)\n",
    "\n",
    "        _, tree_vec, mol_vec = self.encode(mol_batch)\n",
    "        \n",
    "        tree_mean = self.T_mean(tree_vec)\n",
    "        mol_mean = self.G_mean(mol_vec)\n",
    "        \n",
    "        feature_vec =  torch.cat([tree_mean, mol_mean], dim=1)\n",
    "        \n",
    "        return self.ffn(feature_vec)\n",
    "\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "from mol_tree import MolTree\n",
    "import numpy as np\n",
    "\n",
    "# global SMILES_TO_MOLTREE\n",
    "SMILES_TO_MOLTREE = {}\n",
    "\n",
    "\n",
    "class MoleculeDataset(Dataset):\n",
    "\n",
    "    def __init__(self, data_file, raw_path, SMILES_COLUMN = 'smiles', TARGET_COLUMN = 'logP'):\n",
    "        global SMILES_TO_MOLTREE\n",
    "        \n",
    "        self.data = pd.read_csv(data_file)\n",
    "        data_options = ['train','val','test']\n",
    "        \n",
    "        for option in data_options:\n",
    "            if option in data_file:\n",
    "                broken_smiles  = [x.strip(\"\\r\\n \") for x in open(os.path.join(raw_path,option+'_errs.txt'))] \n",
    "                \n",
    "        self.data = self.data[~self.data[SMILES_COLUMN].isin(broken_smiles)]    [:10]\n",
    "        self.SMILES_COLUMN = SMILES_COLUMN\n",
    "        self.TARGET_COLUMN = TARGET_COLUMN\n",
    "        \n",
    "        for i in range(len(self.data)):\n",
    "            if self.data.iloc[i][SMILES_COLUMN] in SMILES_TO_MOLTREE:\n",
    "                mol_tree = SMILES_TO_MOLTREE[self.data.iloc[i][SMILES_COLUMN]]\n",
    "            else:\n",
    "                mol_tree = MolTree(self.data.iloc[i][SMILES_COLUMN])\n",
    "                SMILES_TO_MOLTREE[self.data.iloc[i][SMILES_COLUMN]] = mol_tree\n",
    "                mol_tree.recover()\n",
    "                mol_tree.assemble()\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        global SMILES_TO_MOLTREE\n",
    "        smiles = self.data.iloc[idx][self.SMILES_COLUMN]\n",
    "        target = self.data.iloc[idx][self.TARGET_COLUMN]\n",
    "        if smiles in SMILES_TO_MOLTREE.keys():\n",
    "            mol_tree = SMILES_TO_MOLTREE[smiles]\n",
    "        else:\n",
    "            mol_tree = MolTree(smiles)\n",
    "            SMILES_TO_MOLTREE[smiles] = mol_tree\n",
    "        return mol_tree, target        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create model and load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mol/anaconda3/envs/jtree/lib/python2.7/site-packages/torch/nn/_reduction.py:49: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
      "  warnings.warn(warning.format(ret))\n"
     ]
    }
   ],
   "source": [
    "model = JTPredict(vocab, hidden_size, latent_size, depth, stereo=stereo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.create_ffn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for param in model.parameters():\n",
    "    if param.dim() == 1:\n",
    "        nn.init.constant_(param, 0)\n",
    "    else:\n",
    "        nn.init.xavier_normal_(param)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding.weight tensor([[ 0.0463,  0.0305,  0.0297,  ..., -0.0191, -0.0122,  0.0035],\n",
      "        [-0.0067, -0.0032, -0.0361,  ...,  0.0068,  0.0034, -0.0009],\n",
      "        [ 0.0500, -0.0327, -0.0463,  ..., -0.0318, -0.0108, -0.0416],\n",
      "        ...,\n",
      "        [ 0.0082, -0.0497, -0.0156,  ...,  0.0029, -0.0335, -0.0398],\n",
      "        [-0.0111,  0.0113,  0.0068,  ..., -0.0324,  0.0331,  0.0036],\n",
      "        [ 0.0023, -0.0292,  0.0553,  ...,  0.0406, -0.0140, -0.0168]])\n",
      "jtnn.W_z.weight tensor([[ 0.0491, -0.0374, -0.0435,  ...,  0.0787,  0.0203,  0.0305],\n",
      "        [ 0.0239, -0.0199,  0.0405,  ...,  0.0026, -0.0310, -0.0128],\n",
      "        [-0.0437,  0.0621, -0.0126,  ..., -0.0813,  0.0064,  0.0219],\n",
      "        ...,\n",
      "        [ 0.0091,  0.0150, -0.0158,  ...,  0.0016, -0.0083, -0.0183],\n",
      "        [ 0.0494, -0.0276, -0.0236,  ...,  0.0127,  0.0195, -0.0483],\n",
      "        [ 0.0411, -0.0312, -0.0080,  ...,  0.0044, -0.0467,  0.0318]])\n",
      "jtnn.W_z.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "jtnn.W_r.weight tensor([[-0.0109, -0.0045,  0.0930,  ...,  0.0688, -0.0156,  0.0250],\n",
      "        [ 0.0181,  0.0471, -0.0226,  ..., -0.0518, -0.0163,  0.0315],\n",
      "        [ 0.0106,  0.0666,  0.0162,  ..., -0.0193,  0.0499,  0.0047],\n",
      "        ...,\n",
      "        [ 0.0361, -0.0116, -0.0444,  ...,  0.0726, -0.0858,  0.0397],\n",
      "        [-0.0617,  0.0489,  0.1055,  ...,  0.0242, -0.0732,  0.0672],\n",
      "        [-0.0788,  0.0759,  0.0058,  ..., -0.0413, -0.0535,  0.0841]])\n",
      "jtnn.U_r.weight tensor([[ 0.1051,  0.0141, -0.0011,  ..., -0.0027,  0.0486, -0.0041],\n",
      "        [-0.0697,  0.0065, -0.0320,  ...,  0.0014,  0.0061,  0.0424],\n",
      "        [ 0.1266,  0.0364,  0.0633,  ...,  0.0134, -0.0108, -0.0232],\n",
      "        ...,\n",
      "        [ 0.0020, -0.0548,  0.1070,  ..., -0.0361,  0.0487,  0.0015],\n",
      "        [-0.0196,  0.0111,  0.0074,  ..., -0.0039,  0.0500, -0.0139],\n",
      "        [-0.0215,  0.0083,  0.0356,  ...,  0.0429,  0.0448,  0.0668]])\n",
      "jtnn.U_r.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "jtnn.W_h.weight tensor([[-0.0037, -0.0717,  0.0057,  ...,  0.0077, -0.0737,  0.0067],\n",
      "        [-0.0547, -0.0169, -0.0334,  ..., -0.0361, -0.0390,  0.0255],\n",
      "        [ 0.0588, -0.0175,  0.0219,  ...,  0.0231, -0.0270,  0.0360],\n",
      "        ...,\n",
      "        [ 0.0459, -0.0114,  0.0020,  ..., -0.0037, -0.0193, -0.0234],\n",
      "        [-0.0140,  0.0583, -0.0100,  ..., -0.0300, -0.0655,  0.0288],\n",
      "        [ 0.0160,  0.0019,  0.0482,  ..., -0.0488,  0.0759,  0.0036]])\n",
      "jtnn.W_h.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "jtnn.W.weight tensor([[-0.0070, -0.0377, -0.0511,  ..., -0.0046, -0.0181,  0.0686],\n",
      "        [ 0.0808,  0.0271, -0.0412,  ..., -0.0597,  0.0352,  0.0393],\n",
      "        [-0.0115,  0.0070,  0.0021,  ...,  0.0310,  0.0013, -0.0281],\n",
      "        ...,\n",
      "        [ 0.0795, -0.0266, -0.0113,  ..., -0.0309, -0.0521,  0.0376],\n",
      "        [ 0.0358, -0.0467,  0.0617,  ...,  0.0296,  0.0104,  0.0213],\n",
      "        [-0.0281, -0.0097, -0.0150,  ...,  0.0072,  0.0695, -0.0120]])\n",
      "jtnn.W.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "mpn.W_i.weight tensor([[ 0.0289, -0.0578,  0.0376,  ...,  0.0216, -0.0156, -0.0245],\n",
      "        [ 0.0011,  0.0164,  0.0971,  ...,  0.0002, -0.0448,  0.0690],\n",
      "        [-0.0664, -0.0268,  0.0851,  ...,  0.0508, -0.0606,  0.0454],\n",
      "        ...,\n",
      "        [ 0.0028, -0.0689,  0.0025,  ...,  0.1269,  0.0295, -0.0429],\n",
      "        [-0.0519, -0.0704, -0.0368,  ..., -0.0570, -0.0553, -0.0712],\n",
      "        [ 0.1464, -0.0783,  0.0504,  ...,  0.0294,  0.0731, -0.0252]])\n",
      "mpn.W_h.weight tensor([[ 0.0121, -0.0062, -0.0098,  ..., -0.0223,  0.0436, -0.0022],\n",
      "        [-0.0578, -0.0327,  0.0035,  ...,  0.0013, -0.1038,  0.0230],\n",
      "        [ 0.0139, -0.0879,  0.0266,  ...,  0.0517, -0.0197, -0.0383],\n",
      "        ...,\n",
      "        [ 0.0627, -0.0281, -0.0612,  ...,  0.0249, -0.0574, -0.0040],\n",
      "        [-0.0832, -0.0036,  0.0200,  ...,  0.0405, -0.0036, -0.0241],\n",
      "        [ 0.0144, -0.0013,  0.0222,  ...,  0.0189, -0.0156, -0.1208]])\n",
      "mpn.W_o.weight tensor([[-0.0218, -0.0203,  0.0387,  ..., -0.0473,  0.0478, -0.0053],\n",
      "        [-0.0020, -0.0248,  0.0031,  ...,  0.0058, -0.0044,  0.0123],\n",
      "        [-0.0972,  0.0064, -0.0467,  ...,  0.0504, -0.0852,  0.0208],\n",
      "        ...,\n",
      "        [-0.0458,  0.0731,  0.0205,  ...,  0.0077,  0.0504, -0.0055],\n",
      "        [ 0.0586, -0.0491, -0.0743,  ...,  0.0404, -0.0552, -0.0083],\n",
      "        [-0.0636, -0.0005,  0.0442,  ..., -0.0070, -0.0485, -0.0258]])\n",
      "mpn.W_o.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])\n",
      "T_mean.weight tensor([[-0.0326, -0.0209, -0.0772,  ...,  0.0448,  0.0497,  0.0503],\n",
      "        [ 0.0242, -0.1461,  0.0488,  ...,  0.0270, -0.0198, -0.0014],\n",
      "        [-0.1191,  0.1500,  0.0395,  ...,  0.0060,  0.0496, -0.0702],\n",
      "        ...,\n",
      "        [-0.0730, -0.0153, -0.0786,  ..., -0.1328, -0.1450, -0.0160],\n",
      "        [-0.0251,  0.0445, -0.0934,  ...,  0.0800,  0.0207, -0.0815],\n",
      "        [-0.0110,  0.0121,  0.0424,  ..., -0.0203,  0.0266, -0.1063]])\n",
      "T_mean.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0.])\n",
      "T_var.weight tensor([[-0.0276,  0.0031,  0.0784,  ..., -0.0602,  0.0186, -0.0493],\n",
      "        [ 0.0782,  0.0456, -0.0306,  ..., -0.0319, -0.0282,  0.0821],\n",
      "        [-0.0468, -0.0004, -0.0049,  ..., -0.0859, -0.0505,  0.0265],\n",
      "        ...,\n",
      "        [-0.0359, -0.0079, -0.0114,  ..., -0.0881,  0.0196, -0.0795],\n",
      "        [ 0.0848,  0.0748,  0.0434,  ..., -0.0111, -0.1146,  0.0117],\n",
      "        [ 0.1082,  0.0570, -0.1196,  ...,  0.0256, -0.1680,  0.0175]])\n",
      "T_var.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0.])\n",
      "G_mean.weight tensor([[ 0.0547,  0.0164,  0.0640,  ...,  0.0252, -0.0084,  0.0575],\n",
      "        [-0.0304, -0.0019, -0.1092,  ..., -0.0129,  0.0474, -0.1297],\n",
      "        [-0.1378, -0.0748, -0.0195,  ..., -0.0015, -0.0145,  0.0081],\n",
      "        ...,\n",
      "        [ 0.0525, -0.0534,  0.0083,  ...,  0.0167, -0.0185,  0.1080],\n",
      "        [ 0.0036,  0.1035, -0.0588,  ..., -0.0618, -0.0351, -0.0482],\n",
      "        [ 0.0174, -0.0229, -0.1239,  ..., -0.0543, -0.0728,  0.0940]])\n",
      "G_mean.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0.])\n",
      "G_var.weight tensor([[ 0.0213, -0.0411,  0.0359,  ...,  0.0356, -0.0146,  0.0446],\n",
      "        [-0.0081, -0.0241, -0.0915,  ..., -0.0950, -0.0393, -0.0661],\n",
      "        [-0.0693,  0.0476,  0.0348,  ...,  0.0691, -0.0262,  0.0341],\n",
      "        ...,\n",
      "        [-0.1094,  0.0227, -0.0354,  ...,  0.0125, -0.1388, -0.0164],\n",
      "        [-0.0009,  0.0457,  0.0070,  ...,  0.0990,  0.0809,  0.0071],\n",
      "        [ 0.0613, -0.1259, -0.0274,  ..., -0.0326,  0.0940, -0.0650]])\n",
      "G_var.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0.])\n",
      "ffn.1.weight tensor([[-0.0906, -0.3133,  0.1303,  ..., -0.1678, -0.2319, -0.2223],\n",
      "        [ 0.0267,  0.0584,  0.0042,  ..., -0.0318,  0.2302, -0.1829],\n",
      "        [ 0.0468, -0.0302, -0.2179,  ...,  0.0975,  0.0639,  0.0909],\n",
      "        ...,\n",
      "        [ 0.0759, -0.1012, -0.1698,  ..., -0.0617, -0.1279,  0.2008],\n",
      "        [ 0.0345, -0.0945,  0.0157,  ...,  0.0509,  0.0936,  0.1389],\n",
      "        [-0.1152,  0.0739, -0.0028,  ...,  0.1512, -0.1347, -0.1339]])\n",
      "ffn.1.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.])\n",
      "ffn.4.weight tensor([[ 0.1104, -0.0925,  0.0434,  ...,  0.1937, -0.0719,  0.0526],\n",
      "        [ 0.0602, -0.1469, -0.1582,  ..., -0.0870,  0.0824,  0.2012],\n",
      "        [-0.1320,  0.1234,  0.1687,  ..., -0.1030, -0.0078, -0.0378],\n",
      "        ...,\n",
      "        [ 0.1929,  0.1886,  0.0030,  ..., -0.0433, -0.1526,  0.0185],\n",
      "        [ 0.1433,  0.0100,  0.3022,  ...,  0.0065, -0.0185, -0.0090],\n",
      "        [-0.0837,  0.1812, -0.0791,  ...,  0.0365, -0.2549,  0.2905]])\n",
      "ffn.4.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.])\n",
      "ffn.7.weight tensor([[-0.1975,  0.1618, -0.2512, -0.2361,  0.0810, -0.0652, -0.2099,  0.3093,\n",
      "          0.3060,  0.2265,  0.1251, -0.1239,  0.1891,  0.1764,  0.3673,  0.1286,\n",
      "          0.2934, -0.3947,  0.0574, -0.5267,  0.2699, -0.2439, -0.2550,  0.2081,\n",
      "          0.1409, -0.1278, -0.0208,  0.2078,  0.1396, -0.0285,  0.3486, -0.1499,\n",
      "          0.2901,  0.0223, -0.0744, -0.0607,  0.0075,  0.1774, -0.2489, -0.3042,\n",
      "          0.1305,  0.2034,  0.1850, -0.0236,  0.0462, -0.1273,  0.1803, -0.1508,\n",
      "          0.0682, -0.0702]])\n",
      "ffn.7.bias tensor([0.])\n"
     ]
    }
   ],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print name, param.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load pretrained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jtnn_vae import JTNNVAE\n",
    "\n",
    "model_VAE = JTNNVAE(vocab, hidden_size, latent_size, depth, stereo=stereo)\n",
    "model_VAE.load_state_dict(torch.load(os.path.join(MODEL_PATH, 'MPNVAE-h450-L56-d3-beta0.001/model.iter-4')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.jtnn = model_VAE.jtnn\n",
    "model.mpn = model_VAE.mpn\n",
    "model.embedding = model_VAE.embedding\n",
    "model.T_mean = model_VAE.T_mean\n",
    "model.T_var = model_VAE.T_var\n",
    "model.G_mean = model_VAE.G_mean\n",
    "model.G_var = model_VAE.G_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "del model_VAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding.weight tensor([[ 0.0594,  0.0186, -0.0397,  ...,  0.0985,  0.1281, -0.0242],\n",
      "        [ 0.0607, -0.2044, -0.0376,  ..., -0.0056,  0.0935, -0.0021],\n",
      "        [ 0.0100, -0.0274, -0.0751,  ...,  0.0867, -0.0365,  0.0391],\n",
      "        ...,\n",
      "        [ 0.0686,  0.0343,  0.0617,  ..., -0.0690,  0.0751,  0.0362],\n",
      "        [ 0.0369,  0.0062,  0.0455,  ..., -0.0077, -0.0310, -0.0112],\n",
      "        [ 0.1684,  0.0122, -0.0144,  ...,  0.0881,  0.0752,  0.0060]])\n",
      "jtnn.W_z.weight tensor([[-0.0506,  0.0644,  0.2186,  ..., -0.1192, -0.0397, -0.0724],\n",
      "        [-0.3613,  0.0557,  0.3665,  ...,  0.0265, -0.0641,  0.2011],\n",
      "        [ 0.2609,  0.0762,  0.1762,  ...,  0.2432, -0.1284, -0.0840],\n",
      "        ...,\n",
      "        [ 0.1006, -0.0124,  0.2878,  ...,  0.0095, -0.2124,  0.0286],\n",
      "        [-0.1054, -0.0414,  0.3038,  ...,  0.0535, -0.0193, -0.1250],\n",
      "        [ 0.1979,  0.1238,  0.3254,  ...,  0.0086, -0.0346,  0.4131]])\n",
      "jtnn.W_z.bias tensor([-0.3559, -0.4492, -0.1669, -0.2501, -0.1203, -0.6630, -0.4290, -0.5225,\n",
      "        -0.2735, -0.2099, -0.2350, -0.3181, -0.4130, -0.6144, -0.2755, -0.1980,\n",
      "        -0.3829, -0.1386, -0.4552, -0.3491, -0.2486, -0.5535, -0.2032, -0.0334,\n",
      "        -0.6541, -0.4207, -0.3242, -0.2213, -0.2296, -0.2190, -0.4415, -0.4160,\n",
      "         0.0759, -0.4647, -0.1276, -0.2763, -0.6552, -0.1309, -0.3283, -0.3418,\n",
      "        -0.2228, -0.3096, -0.3210, -0.2496, -0.2437, -0.3303, -0.5993, -0.4107,\n",
      "        -0.1627, -0.1187, -0.3123, -0.3802, -0.5888, -0.5393, -0.1198, -0.2805,\n",
      "        -0.3724, -0.5290, -0.5146, -0.0449, -0.4403, -0.4816, -0.4376, -0.3012,\n",
      "        -0.0628, -0.2575, -0.3405, -0.1674, -0.3517, -0.2573, -0.1473, -0.2265,\n",
      "        -0.3785, -0.4330, -0.4093, -0.4065, -0.4419, -0.3123, -0.0230, -0.1471,\n",
      "        -0.1474, -0.3027, -0.3608, -0.3887, -0.5712, -0.2256, -0.0579,  0.1302,\n",
      "        -0.3445, -0.2621, -0.1601, -0.4146, -0.4147, -0.5464, -0.4028, -0.3980,\n",
      "        -0.3997, -0.3359, -0.5922, -0.3560, -0.3341, -0.3395, -0.4873, -0.5217,\n",
      "        -0.3387, -0.5342, -0.4966, -0.1405, -0.4012, -0.1456, -0.2769, -0.5275,\n",
      "         0.0417, -0.2511, -0.2483, -0.3845, -0.3275, -0.1463,  0.2401, -0.1542,\n",
      "        -0.0599, -0.2106, -0.2672, -0.4294, -0.3189, -0.4623, -0.3365, -0.1230,\n",
      "         0.0258, -0.4149, -0.2094, -0.1694, -0.0221, -0.3393, -0.0937, -0.4972,\n",
      "        -0.3490, -0.4198, -0.4348, -0.4950, -0.1664, -0.1990, -0.5436, -0.2671,\n",
      "        -0.3991, -0.4401, -0.4115, -0.3858, -0.2091, -0.3786, -0.3904, -0.4741,\n",
      "        -0.2898, -0.3414, -0.4786, -0.5933, -0.2843, -0.3682, -0.2954, -0.2927,\n",
      "        -0.5679, -0.3984, -0.1343, -0.1511, -0.5054, -0.3564, -0.4148, -0.3438,\n",
      "        -0.2294, -0.0457, -0.4349, -0.6062, -0.2143,  0.0033, -0.7105, -0.4107,\n",
      "        -0.1968, -0.2075, -0.3761, -0.0850, -0.3380, -0.1349, -0.1066, -0.4101,\n",
      "        -0.4917,  0.0524, -0.2362, -0.3249, -0.4849, -0.0387, -0.5123, -0.3523,\n",
      "        -0.4378, -0.2433, -0.3962, -0.3615, -0.5628, -0.0094, -0.4181, -0.2874,\n",
      "        -0.1124, -0.4116, -0.1859, -0.3292, -0.2452, -0.4107, -0.4839, -0.5222,\n",
      "        -0.3409, -0.4384, -0.3778, -0.0999, -0.4587, -0.3137, -0.2192, -0.2925,\n",
      "        -0.4676, -0.4242, -0.3036,  0.0878, -0.1922, -0.1245, -0.1879, -0.3584,\n",
      "        -0.0513, -0.1922, -0.3470,  0.1219, -0.5235, -0.2737, -0.3486, -0.1352,\n",
      "        -0.3159, -0.3198, -0.3720, -0.2038, -0.3149, -0.3379, -0.2673, -0.3052,\n",
      "        -0.2175, -0.3460, -0.5626, -0.6118, -0.4867, -0.3665, -0.0977,  0.0311,\n",
      "        -0.4426, -0.2781, -0.4256,  0.0224, -0.0038, -0.3784, -0.3574, -0.5144,\n",
      "        -0.3601, -0.6224, -0.3419, -0.1706, -0.3075, -0.1505, -0.4784, -0.2971,\n",
      "        -0.3232, -0.1716, -0.3105, -0.4709, -0.0860, -0.2583, -0.3751, -0.3932,\n",
      "        -0.4161, -0.6287, -0.3957, -0.4370, -0.4946,  0.1248, -0.2031, -0.2135,\n",
      "        -0.2498, -0.3637, -0.1728, -0.2380, -0.1634, -0.2282, -0.2975, -0.5295,\n",
      "        -0.1784, -0.3450, -0.3946, -0.0652, -0.4920, -0.5609, -0.4953, -0.3280,\n",
      "        -0.3550, -0.4559, -0.3358, -0.0254, -0.2960, -0.4160, -0.5834, -0.4080,\n",
      "        -0.2742, -0.0955, -0.4527, -0.6236, -0.5778, -0.2529, -0.2005, -0.2865,\n",
      "        -0.5522, -0.0481, -0.4884, -0.5994, -0.3561, -0.3930, -0.0233, -0.1389,\n",
      "        -0.6273, -0.4689, -0.3383, -0.3028, -0.3131, -0.3684, -0.3964,  0.1258,\n",
      "        -0.4634, -0.3105, -0.1249, -0.0634, -0.1846, -0.3480, -0.4101, -0.0438,\n",
      "        -0.1688, -0.4116, -0.4263, -0.0439, -0.3156, -0.3765, -0.2597, -0.0595,\n",
      "        -0.3477, -0.5448, -0.0957, -0.3034, -0.2975, -0.4000, -0.2306, -0.2663,\n",
      "        -0.2266, -0.3456, -0.3325, -0.6662, -0.3485, -0.1270, -0.5654, -0.2454,\n",
      "        -0.1511, -0.3810, -0.1449, -0.4561, -0.2876,  0.2440, -0.3702, -0.2124,\n",
      "        -0.2981, -0.3560, -0.2366, -0.3872, -0.3362, -0.0909, -0.4276, -0.5786,\n",
      "        -0.2877, -0.0516, -0.4632, -0.5040, -0.4524, -0.4149, -0.4918, -0.2605,\n",
      "        -0.4214, -0.3123, -0.4841, -0.3940,  0.2032, -0.3815,  0.0631, -0.2416,\n",
      "        -0.5356, -0.4468, -0.3047, -0.4494, -0.4881, -0.1914, -0.1754, -0.4021,\n",
      "        -0.3038, -0.3890, -0.3099, -0.4798, -0.4707, -0.4584, -0.4799, -0.4082,\n",
      "        -0.3020, -0.1916, -0.4491, -0.1644, -0.1792, -0.2288, -0.0751, -0.3997,\n",
      "        -0.1637, -0.2966, -0.2951, -0.2813, -0.0960, -0.4636, -0.2481, -0.4379,\n",
      "        -0.6527, -0.1767, -0.0836, -0.3205, -0.1970, -0.4290, -0.5267, -0.1250,\n",
      "        -0.3642, -0.4649, -0.2737, -0.5629, -0.4818, -0.3806, -0.3529, -0.3801,\n",
      "        -0.4474, -0.2774, -0.3590, -0.3161, -0.4458,  0.3128, -0.2139, -0.4822,\n",
      "        -0.3738, -0.3667])\n",
      "jtnn.W_r.weight tensor([[ 0.0173, -0.0078,  0.1085,  ..., -0.0135, -0.0339, -0.2878],\n",
      "        [ 0.0299, -0.0616,  0.0348,  ...,  0.0731, -0.0784, -0.0533],\n",
      "        [ 0.0131,  0.2780, -0.0145,  ..., -0.0065, -0.2691, -0.3528],\n",
      "        ...,\n",
      "        [-0.2006, -0.0197,  0.1802,  ..., -0.1983, -0.4799, -0.4469],\n",
      "        [-0.0156, -0.2799,  0.0203,  ...,  0.0204, -0.2131, -0.0182],\n",
      "        [-0.0501, -0.0349,  0.2719,  ...,  0.0098, -0.3342, -0.2933]])\n",
      "jtnn.U_r.weight tensor([[-0.1307,  0.1392,  0.0710,  ...,  0.1587,  0.0364, -0.1776],\n",
      "        [ 0.1390, -0.0046,  0.0307,  ...,  0.0184,  0.3096, -0.2144],\n",
      "        [-0.0743, -0.0315, -0.0194,  ...,  0.1324, -0.0690, -0.0437],\n",
      "        ...,\n",
      "        [-0.0232,  0.0910,  0.1173,  ...,  0.0101,  0.1067,  0.2811],\n",
      "        [-0.1891, -0.0192, -0.1469,  ...,  0.0474, -0.0733,  0.0191],\n",
      "        [-0.0721, -0.1624, -0.0805,  ...,  0.0030,  0.0942,  0.0378]])\n",
      "jtnn.U_r.bias tensor([-1.2035e-01, -9.4693e-02, -6.5235e-02, -2.3574e-01, -1.2522e-01,\n",
      "        -1.1945e-01, -2.2198e-01, -8.0021e-02, -1.0293e-01, -6.4858e-02,\n",
      "        -2.6512e-02, -1.0687e-01, -7.2757e-02, -1.3835e-01, -1.5047e-01,\n",
      "        -1.2551e-01, -1.3728e-01, -1.2846e-01, -1.4413e-01, -3.1500e-03,\n",
      "        -8.1714e-02, -1.3777e-01, -9.2818e-02, -1.0002e-01, -1.0365e-01,\n",
      "        -1.6451e-01, -1.0772e-01, -1.0449e-01, -1.2027e-01, -7.0041e-02,\n",
      "        -1.5054e-01, -1.5160e-01, -5.6584e-02, -1.4677e-01, -8.3488e-02,\n",
      "        -1.2902e-02, -1.3188e-01, -9.2793e-02, -1.0034e-01, -7.7746e-02,\n",
      "        -1.4707e-01, -7.4266e-02, -2.4252e-02, -4.1697e-02,  1.9104e-02,\n",
      "        -9.2459e-02, -1.8623e-01, -1.3568e-01, -1.0883e-01, -9.5694e-03,\n",
      "        -4.1417e-02, -1.3077e-01, -1.8451e-01, -8.5946e-03, -1.8001e-02,\n",
      "        -1.2463e-01, -1.0110e-01, -1.5769e-01, -7.7532e-02, -8.9746e-02,\n",
      "        -7.7386e-02, -7.6046e-02, -6.8969e-02, -1.3123e-01,  1.3043e-02,\n",
      "        -1.3457e-01,  2.3377e-02, -9.5480e-02, -1.8132e-01, -1.4841e-01,\n",
      "        -8.8586e-02, -1.5224e-01, -1.3979e-01, -1.0971e-01, -2.9726e-02,\n",
      "        -3.1071e-02, -1.3373e-01, -8.6295e-02, -1.7911e-02,  4.8859e-02,\n",
      "        -7.8023e-02, -8.2147e-02, -1.4346e-01, -4.5232e-02, -1.7562e-01,\n",
      "        -7.4752e-02, -1.0622e-01,  9.0045e-03, -7.4093e-02, -1.2468e-01,\n",
      "        -1.0901e-01, -9.7934e-02, -7.6155e-02, -1.6826e-01, -1.0117e-01,\n",
      "        -6.6205e-02, -1.8440e-01, -1.3011e-01, -8.3491e-02, -5.9971e-02,\n",
      "        -1.0664e-01, -7.5390e-02, -1.1628e-01, -1.7220e-01, -4.5119e-02,\n",
      "        -3.9446e-02, -1.3910e-01, -1.2813e-01, -9.1147e-02, -1.5656e-02,\n",
      "        -1.0042e-01, -3.0993e-02, -6.1432e-02, -9.8221e-02, -9.8164e-02,\n",
      "        -1.1648e-01, -1.3210e-01, -1.3455e-01, -8.4769e-02, -1.2145e-01,\n",
      "        -1.3139e-01, -4.8285e-02, -5.4467e-02, -8.3011e-02, -6.6262e-02,\n",
      "        -8.2084e-02, -8.8515e-02, -7.4658e-02,  5.3365e-04, -5.9378e-02,\n",
      "        -7.0663e-02, -1.1177e-01, -7.2681e-02,  2.7974e-03, -1.2041e-02,\n",
      "        -1.7319e-01, -1.7131e-01, -1.1285e-01, -1.6424e-01, -5.4209e-02,\n",
      "        -7.3449e-02, -1.4863e-02, -1.6037e-01, -7.9793e-02,  1.1300e-03,\n",
      "        -1.1836e-01, -6.2546e-02, -1.3921e-01, -7.9395e-02, -1.7093e-01,\n",
      "        -1.2101e-01, -1.0568e-01, -2.0189e-01, -1.2737e-02, -1.6197e-01,\n",
      "        -2.1711e-01, -1.2923e-01, -5.8128e-02, -4.7292e-02, -4.1531e-02,\n",
      "        -1.3371e-01, -2.0357e-01,  1.3538e-02, -3.2231e-02, -1.6059e-01,\n",
      "        -1.3738e-01, -7.8121e-02, -8.5222e-02, -1.3091e-01, -1.3596e-01,\n",
      "        -1.0578e-01, -1.2214e-01, -7.9522e-02, -4.0335e-02, -1.5530e-01,\n",
      "        -7.7036e-02, -5.1600e-02, -4.1564e-02, -8.0504e-02, -2.7772e-02,\n",
      "        -1.0861e-01, -8.8764e-02, -1.0032e-01, -1.0378e-01, -9.6383e-02,\n",
      "        -1.6817e-02, -2.4832e-02, -1.6058e-01, -3.7526e-02, -6.2179e-02,\n",
      "        -9.9092e-02, -5.1042e-02, -9.4242e-02, -1.0149e-01, -2.7023e-02,\n",
      "        -2.9136e-01, -1.2032e-01, -1.7323e-01, -2.0102e-01, -6.3670e-02,\n",
      "        -1.5495e-01, -9.5399e-02, -3.1912e-02, -1.0016e-01, -8.9161e-02,\n",
      "        -1.6056e-01, -7.8475e-02,  2.4597e-02, -1.2112e-01, -9.3080e-02,\n",
      "        -1.1491e-01, -5.6928e-02, -5.9567e-02, -1.4897e-02, -1.2102e-01,\n",
      "        -8.7761e-02, -1.1144e-01, -8.8830e-02, -4.8150e-02, -7.4321e-02,\n",
      "        -8.9158e-02, -6.2726e-02, -3.1403e-02, -6.3730e-02, -4.5539e-02,\n",
      "        -5.2818e-02, -6.8782e-02, -1.1217e-01, -9.1220e-02, -7.5787e-02,\n",
      "        -3.2242e-02, -1.3683e-02, -9.5851e-02, -7.0333e-02, -1.1799e-01,\n",
      "        -1.3931e-01, -8.6022e-02, -6.2211e-02, -1.5724e-01, -9.7825e-02,\n",
      "        -2.6607e-04, -7.2627e-02, -1.2832e-01, -1.6311e-01, -7.9975e-02,\n",
      "        -9.2650e-02, -9.3715e-02, -6.6729e-03, -7.6282e-02, -7.9108e-02,\n",
      "        -6.5212e-02, -1.3012e-01, -6.4427e-02, -1.4857e-01, -1.0656e-01,\n",
      "        -2.9086e-01, -7.8752e-02, -8.4237e-02, -1.4697e-01, -8.9377e-02,\n",
      "        -2.8481e-02, -6.2614e-02, -8.6508e-03, -7.7967e-02, -1.4224e-01,\n",
      "        -7.2650e-02, -8.3791e-02, -9.7100e-02, -1.0511e-01, -9.9012e-02,\n",
      "        -1.4295e-01, -6.9937e-03, -7.4619e-02, -1.6157e-01, -3.9655e-02,\n",
      "        -9.2164e-02, -1.2232e-01, -5.6692e-02, -1.3824e-01, -1.0746e-01,\n",
      "        -1.3031e-01, -1.0168e-01,  1.0983e-02, -2.2710e-02, -1.9923e-01,\n",
      "        -9.5404e-02, -8.4317e-02, -3.3281e-02, -8.0546e-02, -1.4056e-01,\n",
      "        -8.3173e-02, -5.1234e-02, -1.9794e-01, -1.3828e-01, -1.2833e-01,\n",
      "        -1.3076e-01, -4.0222e-02, -1.2493e-01, -2.3921e-02, -8.3368e-02,\n",
      "        -9.9057e-02, -1.4211e-01, -6.9581e-02, -8.1331e-02, -7.0587e-02,\n",
      "        -8.2842e-02, -1.1987e-01, -8.6102e-02, -1.0097e-01, -3.9226e-02,\n",
      "        -9.5271e-02, -3.1884e-01, -1.2684e-01, -6.0456e-02, -1.1591e-01,\n",
      "        -1.2922e-01, -8.5517e-02, -1.3459e-01, -3.9834e-02, -3.4470e-02,\n",
      "        -1.8415e-01, -1.3063e-01, -5.0865e-02, -7.2012e-02, -2.9145e-02,\n",
      "        -7.2474e-02, -2.1218e-01, -4.0204e-02, -1.8152e-01, -5.1762e-02,\n",
      "        -3.1819e-02, -8.1585e-02, -1.6128e-01, -1.3351e-01, -1.0218e-01,\n",
      "        -1.1820e-01, -1.5760e-02, -9.0270e-02, -2.7753e-02, -2.6651e-02,\n",
      "        -1.0885e-01, -1.4100e-01, -1.0658e-01, -9.1814e-02, -9.9759e-02,\n",
      "        -9.4224e-02, -9.7049e-02, -5.0912e-02, -6.6454e-02, -1.1325e-01,\n",
      "        -8.0655e-02, -1.0204e-01, -1.0709e-01, -2.8330e-02, -1.1249e-01,\n",
      "        -1.5848e-01, -8.2992e-02, -4.0373e-02, -1.1427e-01, -9.1253e-02,\n",
      "        -6.6536e-02, -1.0068e-01, -5.4385e-02, -1.2515e-01, -6.7386e-02,\n",
      "        -1.6820e-01, -1.1459e-01, -1.0832e-02, -3.0741e-02,  2.4687e-03,\n",
      "        -7.7963e-02, -1.1533e-01, -5.7848e-02, -4.7454e-02, -9.6388e-02,\n",
      "        -1.7762e-01, -1.1210e-01, -9.9417e-02, -9.0538e-02, -9.2928e-02,\n",
      "        -7.1573e-02, -1.3186e-01, -1.0312e-01, -5.1429e-02, -1.4147e-01,\n",
      "        -9.5948e-02, -1.3259e-01, -7.3100e-02, -1.1502e-01, -8.3720e-02,\n",
      "        -2.4359e-02, -6.5563e-02, -4.4095e-02, -1.1608e-01, -6.6253e-02,\n",
      "        -8.8990e-02, -4.4083e-02,  2.5923e-02, -8.1973e-02, -3.0802e-02,\n",
      "        -2.3823e-01, -2.6320e-02, -1.0412e-01, -1.0978e-01, -2.0860e-01,\n",
      "        -7.6999e-02, -9.2554e-02, -3.7526e-02, -6.7190e-02, -1.8288e-01,\n",
      "        -3.5868e-02, -1.4074e-01, -2.0783e-01, -1.1623e-01, -1.1634e-01,\n",
      "        -1.1960e-01, -1.1683e-01, -5.0591e-02, -7.2599e-02, -9.4390e-02,\n",
      "        -9.2903e-02, -1.0657e-01, -7.3530e-02, -6.4884e-02, -1.4740e-01,\n",
      "        -1.3382e-01, -7.2709e-02, -9.4304e-02, -7.9604e-02, -1.2231e-01,\n",
      "        -1.2269e-01, -3.6239e-02, -8.8182e-02, -1.9115e-01, -1.1065e-01,\n",
      "        -1.6741e-01, -8.0564e-02, -6.5577e-02, -1.1554e-01, -1.2106e-01,\n",
      "        -1.2517e-01, -2.8832e-02, -1.3204e-01, -5.2679e-02, -1.0708e-01,\n",
      "         1.1330e-02, -3.5753e-02, -1.5977e-01, -7.8739e-02, -1.1626e-01])\n",
      "jtnn.W_h.weight tensor([[-3.8648e-02, -1.9340e-01,  1.5817e-02,  ...,  1.6628e-01,\n",
      "          1.4871e-02,  2.3540e-02],\n",
      "        [ 3.7299e-02, -8.7448e-02,  1.3369e-01,  ...,  1.0440e-01,\n",
      "          5.6763e-02,  9.6660e-02],\n",
      "        [-1.4984e-01, -9.7230e-02,  1.3612e-01,  ..., -2.3928e-05,\n",
      "          1.2969e-01,  2.6117e-02],\n",
      "        ...,\n",
      "        [ 2.9593e-02,  1.0796e-01,  2.9542e-02,  ..., -5.8565e-02,\n",
      "         -1.3677e-01,  2.0764e-02],\n",
      "        [-3.1683e-02,  7.7659e-02, -3.8658e-02,  ...,  3.9518e-02,\n",
      "         -3.1690e-01,  5.4457e-02],\n",
      "        [ 1.4454e-05,  7.0564e-02, -1.4426e-02,  ..., -4.2586e-02,\n",
      "          9.2897e-02, -2.3553e-01]])\n",
      "jtnn.W_h.bias tensor([-6.3721e-03, -1.6656e-02, -3.1830e-02, -1.4079e-02,  3.8714e-02,\n",
      "         9.6908e-03,  1.6919e-03, -7.0509e-04, -1.4563e-02, -5.0628e-02,\n",
      "         2.8071e-02,  3.4536e-03,  5.6235e-03,  4.3815e-02,  2.1464e-02,\n",
      "         1.6730e-02,  5.9917e-02, -7.8196e-02, -8.5849e-03, -4.2485e-02,\n",
      "        -3.8583e-02,  1.1941e-02,  2.4389e-02, -4.8315e-03,  1.4384e-03,\n",
      "        -1.7062e-02, -7.8519e-03,  2.0880e-02, -2.7086e-02,  5.0613e-03,\n",
      "         4.3289e-03,  2.2424e-02,  8.7238e-02, -4.1139e-02,  1.5465e-02,\n",
      "         5.5807e-03, -5.3496e-04, -5.9865e-02,  1.6919e-02, -3.4730e-02,\n",
      "        -2.2511e-02, -2.3658e-02,  6.1290e-03,  1.3267e-02, -4.0347e-03,\n",
      "        -1.8969e-03,  1.0096e-02,  1.4720e-02, -3.8413e-05,  4.3011e-02,\n",
      "         2.8646e-02,  6.6117e-03, -5.9623e-03,  7.1820e-04, -5.5428e-02,\n",
      "         5.4163e-02,  4.9377e-02, -6.0411e-02, -5.0688e-02, -8.8837e-02,\n",
      "         2.7906e-02,  2.1733e-02,  8.9886e-03,  1.2117e-04, -7.5643e-02,\n",
      "         4.8157e-03,  8.2689e-02,  8.5097e-02, -5.3784e-02,  2.1244e-02,\n",
      "         4.6975e-02,  2.8007e-02, -1.9673e-02,  1.2934e-02, -6.8571e-03,\n",
      "        -3.5301e-02, -2.1917e-02,  3.8351e-02, -3.7240e-02, -6.9281e-02,\n",
      "         1.1435e-02, -3.2775e-02,  1.8445e-02, -4.9450e-02, -8.0175e-03,\n",
      "        -5.3756e-02, -5.4101e-02,  8.9003e-02,  1.5688e-02, -1.7897e-02,\n",
      "        -4.7598e-03, -2.3231e-02,  2.3707e-03, -3.7440e-02,  2.3804e-02,\n",
      "         7.6943e-03, -5.2635e-02,  5.2862e-02, -4.4969e-02, -5.0609e-02,\n",
      "         1.5304e-02, -2.6317e-02,  1.3008e-02, -2.0004e-02,  2.1807e-02,\n",
      "         3.5401e-03,  2.6471e-02,  1.3696e-02,  2.0792e-02, -4.1480e-02,\n",
      "        -3.3343e-02, -1.5790e-02, -9.2713e-02,  7.2064e-02, -2.3887e-03,\n",
      "        -1.4077e-02,  4.5266e-03, -3.0156e-02,  1.1146e-01, -2.2583e-02,\n",
      "         6.0990e-02,  2.0057e-02, -4.2100e-02,  1.6518e-02,  3.8890e-02,\n",
      "         2.5995e-03,  1.5130e-02, -5.7878e-02, -1.4066e-01,  5.7187e-03,\n",
      "        -1.3483e-01, -2.7338e-03,  3.6249e-02, -1.6520e-02, -1.2828e-01,\n",
      "        -2.5446e-02, -5.7840e-02, -2.7355e-02, -4.0122e-02, -8.5543e-03,\n",
      "         2.9981e-02, -8.9518e-03, -1.3434e-02, -2.0896e-02, -3.8476e-02,\n",
      "         1.6705e-02, -8.0021e-02, -6.5165e-03, -9.6736e-03,  4.9834e-02,\n",
      "         1.9791e-02,  4.9150e-03, -4.7793e-03,  1.3641e-02,  1.4087e-02,\n",
      "         6.4640e-03,  4.2172e-02,  5.0587e-03, -6.4484e-04,  6.1630e-02,\n",
      "         4.2838e-02, -1.7497e-02, -6.7262e-02, -7.5586e-03, -6.7317e-03,\n",
      "         2.0829e-02, -1.6586e-03, -4.9835e-02, -4.0124e-02, -2.4405e-02,\n",
      "         5.7072e-02,  1.7471e-03,  3.3770e-02,  8.0646e-02,  1.8510e-02,\n",
      "         2.3029e-02,  7.3852e-02,  1.5105e-02, -3.5234e-02,  1.5236e-01,\n",
      "        -2.6239e-02,  4.3793e-02,  4.2846e-02,  8.6671e-03, -1.2395e-03,\n",
      "         9.6963e-02,  8.7038e-04,  4.3967e-02, -3.7724e-02,  7.7350e-03,\n",
      "         3.3266e-03,  1.6962e-02,  4.6200e-04, -3.2544e-02,  7.7462e-03,\n",
      "        -4.3522e-02,  1.3551e-02,  6.0389e-02, -7.7832e-03,  7.0366e-02,\n",
      "         2.8795e-02, -2.9578e-03, -1.3051e-01, -2.2164e-02,  2.6623e-02,\n",
      "        -3.5094e-02,  1.8220e-02, -3.0053e-02, -2.1087e-02, -5.0596e-02,\n",
      "        -3.1380e-02, -9.8016e-02, -3.0767e-02, -7.2893e-04, -9.3601e-03,\n",
      "         3.9903e-03,  1.3268e-02,  3.1524e-04, -2.9888e-02, -3.2838e-02,\n",
      "        -5.5844e-02, -1.2701e-01, -5.1978e-03, -1.2941e-02,  1.0636e-01,\n",
      "         6.0213e-02,  1.1921e-03, -2.4521e-02,  1.5147e-02, -2.4567e-02,\n",
      "        -3.8951e-02, -6.1635e-02,  6.0548e-02, -7.3838e-02,  6.9281e-03,\n",
      "         3.6618e-02, -3.4695e-02,  2.8980e-02,  7.0305e-02, -1.1861e-04,\n",
      "         2.6109e-02,  2.3611e-02, -2.3362e-02,  4.2581e-02, -3.1864e-02,\n",
      "        -4.7766e-02,  6.1974e-02, -8.8798e-02,  1.9200e-02, -9.1018e-03,\n",
      "        -3.0598e-02,  2.6335e-03, -7.7853e-02,  4.1376e-02,  2.6974e-02,\n",
      "        -2.5901e-02, -6.2709e-02, -1.2175e-02,  2.6852e-02, -5.5895e-02,\n",
      "         7.0470e-02, -6.6090e-02,  1.4155e-02,  2.2288e-02,  1.3909e-02,\n",
      "        -3.0038e-02,  3.9415e-02,  3.4271e-02,  4.6955e-02,  3.2939e-02,\n",
      "        -3.8031e-02,  4.1068e-02,  2.5897e-02, -5.4996e-02,  7.5841e-03,\n",
      "        -4.4076e-02,  2.0427e-02, -1.3131e-01, -9.0476e-02,  4.1963e-02,\n",
      "        -2.8961e-02, -2.3802e-02, -3.4025e-02,  3.7733e-02, -2.0405e-02,\n",
      "        -4.5901e-02,  9.5612e-03, -2.2240e-02,  9.2267e-03, -2.9879e-02,\n",
      "         1.3536e-02, -6.0494e-02, -3.2422e-02,  2.4549e-02,  1.5918e-02,\n",
      "         3.1543e-02, -5.1348e-02, -3.4207e-02,  1.1063e-02,  1.3499e-02,\n",
      "        -1.0448e-02,  4.4611e-03,  2.0589e-02, -2.5838e-02, -5.2750e-03,\n",
      "         6.1570e-02,  1.9109e-02,  8.0946e-04,  1.8848e-02, -4.4957e-02,\n",
      "         2.1531e-02, -2.9848e-03,  3.5269e-04,  1.4629e-01,  4.3083e-02,\n",
      "         1.2693e-02,  5.3310e-03, -2.9066e-02,  1.3281e-02,  3.4214e-02,\n",
      "        -2.0876e-02,  2.0455e-02, -2.6749e-02, -6.8912e-02,  2.2097e-02,\n",
      "        -1.3585e-02, -2.3231e-02, -1.6534e-01, -1.7831e-02, -2.5116e-02,\n",
      "         7.0814e-02, -3.5724e-02, -6.5349e-02,  4.4248e-02, -2.1690e-02,\n",
      "        -1.4107e-02,  3.6527e-02, -1.6950e-02,  2.5917e-02,  7.7333e-02,\n",
      "        -2.1260e-02, -1.2257e-02,  6.0096e-03, -7.0400e-02, -2.1688e-02,\n",
      "        -1.7217e-03,  9.2692e-02,  4.1254e-02,  9.7204e-03, -3.7603e-02,\n",
      "         4.3719e-03, -3.0476e-02,  5.3467e-02,  8.9900e-03,  6.4226e-02,\n",
      "         9.8246e-03, -1.1228e-02,  1.1158e-01, -4.4358e-03, -5.2026e-02,\n",
      "        -6.0427e-02, -2.2889e-03,  2.1580e-02,  1.3552e-02, -9.1543e-02,\n",
      "         6.2350e-02,  4.8111e-02, -2.9203e-03, -5.4057e-02,  2.8170e-02,\n",
      "         4.3175e-02,  1.7713e-02,  3.9165e-02, -1.4976e-02, -1.9371e-02,\n",
      "         3.2554e-04,  4.0616e-02,  4.2861e-02,  1.0241e-02,  1.4970e-02,\n",
      "        -2.3970e-02,  1.9499e-02, -4.9135e-03, -3.0815e-03,  1.4369e-02,\n",
      "        -1.1838e-02, -3.0391e-02, -4.1594e-02,  4.5467e-02,  4.3168e-02,\n",
      "         9.8397e-02, -2.2653e-02,  1.9049e-03,  2.8597e-02, -9.9355e-03,\n",
      "        -1.4767e-02,  5.8068e-03, -2.2242e-02, -2.4465e-02, -2.7603e-02,\n",
      "         8.3438e-02, -1.7855e-02, -6.8659e-02,  2.3395e-02, -7.2630e-04,\n",
      "         1.7505e-02,  1.1045e-02, -1.7463e-02, -2.0050e-02,  1.8732e-02,\n",
      "         2.7894e-03, -1.0136e-02,  3.5879e-02, -6.6521e-02,  3.8869e-02,\n",
      "         4.0912e-03, -8.8993e-03,  1.8911e-02, -9.2295e-02,  6.4503e-02,\n",
      "         2.2393e-02, -2.5471e-02,  9.0416e-02,  2.2955e-02, -2.3983e-03,\n",
      "         2.1131e-02,  9.9647e-03, -2.9273e-02, -9.7095e-02,  3.0223e-02,\n",
      "        -1.2892e-02,  5.5003e-02, -2.8290e-02, -3.7866e-03,  8.4226e-03,\n",
      "        -1.1224e-02, -1.1508e-02, -6.4412e-03, -3.3314e-03, -2.7500e-02,\n",
      "         2.2182e-02, -3.5514e-02,  4.3681e-03,  3.0560e-02, -2.6158e-03,\n",
      "        -1.6803e-01, -1.4784e-02, -1.7015e-02,  5.1898e-03,  5.1311e-02])\n",
      "jtnn.W.weight tensor([[-0.0926,  0.3228,  0.3122,  ...,  0.0424, -0.0546, -0.0072],\n",
      "        [-0.0775, -0.2025, -0.0038,  ..., -0.1501,  0.0959, -0.0431],\n",
      "        [ 0.1871,  0.0458,  0.0216,  ..., -0.0984,  0.0921, -0.1538],\n",
      "        ...,\n",
      "        [-0.0220, -0.1256, -0.0757,  ..., -0.1608, -0.1024,  0.0808],\n",
      "        [ 0.0716,  0.0825, -0.1157,  ..., -0.0007, -0.1175, -0.1451],\n",
      "        [-0.2133,  0.0179, -0.0052,  ...,  0.1055, -0.0047,  0.0072]])\n",
      "jtnn.W.bias tensor([ 6.7460e-03, -5.1538e-02,  2.0651e-02, -4.3035e-02, -7.5362e-02,\n",
      "         4.8616e-02, -5.8535e-03,  4.7824e-02, -2.5815e-02, -3.0010e-02,\n",
      "         3.6183e-02,  1.8894e-02, -1.0431e-01, -1.3781e-01, -9.7522e-02,\n",
      "        -1.7715e-02, -1.0294e-01,  1.3252e-01,  1.6226e-02, -1.0553e-01,\n",
      "         3.4645e-03,  3.3269e-02,  1.4273e-01,  3.2340e-02, -1.6323e-01,\n",
      "         1.6847e-01,  1.0849e-01,  1.8903e-02,  6.6168e-02,  5.6639e-02,\n",
      "        -1.1619e-01,  7.4772e-02, -1.2104e-01,  1.1719e-02, -1.3116e-02,\n",
      "         8.8787e-02, -1.8996e-02,  6.4969e-02, -2.7022e-02,  2.4552e-02,\n",
      "         2.1283e-02,  1.7273e-03, -1.2679e-01,  2.2832e-02,  3.8464e-02,\n",
      "        -7.2474e-02,  7.1235e-03, -8.7922e-03,  1.3735e-02,  2.8266e-02,\n",
      "        -6.2204e-03, -6.3420e-02,  9.9842e-02, -6.1704e-02, -1.0424e-01,\n",
      "         6.9251e-03,  9.7783e-03, -8.9776e-02, -1.5932e-01, -5.1507e-02,\n",
      "        -9.1804e-02, -2.9181e-02, -2.9623e-02, -1.6871e-01,  2.6690e-02,\n",
      "        -1.6940e-01, -1.0488e-02,  1.2333e-01,  1.1619e-01, -4.4223e-02,\n",
      "         1.5838e-01,  3.2806e-03,  2.8091e-02, -5.4333e-02, -6.1062e-02,\n",
      "        -2.5780e-02,  1.6892e-01, -1.1741e-01, -7.3853e-04,  8.5096e-02,\n",
      "         7.3635e-04,  4.5201e-02, -7.0936e-02, -3.9491e-02,  4.9415e-02,\n",
      "        -3.2572e-02,  1.0627e-01, -9.1650e-02, -3.4664e-02, -8.0932e-02,\n",
      "         1.3184e-02, -7.0325e-02, -4.1177e-02, -6.8426e-02,  6.3842e-02,\n",
      "         4.5076e-02, -4.7219e-02,  5.3485e-02, -7.6127e-02, -1.9292e-01,\n",
      "         2.3385e-02, -1.5557e-02, -1.3111e-01,  1.1298e-01, -5.4003e-02,\n",
      "        -6.1631e-02, -2.2091e-02, -1.1059e-01, -4.2845e-02, -1.3521e-02,\n",
      "        -9.3982e-02,  4.0170e-02,  1.5533e-01, -9.5557e-03, -2.7271e-02,\n",
      "         6.4798e-02, -6.6982e-02, -6.8766e-02, -1.2514e-01, -4.3531e-02,\n",
      "        -4.4297e-02, -1.4120e-01, -3.2536e-02, -2.6769e-02, -5.1685e-02,\n",
      "        -5.6545e-02,  6.9446e-02,  4.1495e-02, -5.7073e-02,  9.0850e-02,\n",
      "        -2.2861e-02,  4.2911e-02, -7.9476e-02,  1.2081e-01, -7.4774e-02,\n",
      "        -7.1515e-02,  1.0018e-01, -1.8915e-03,  4.4111e-02, -6.9597e-02,\n",
      "        -5.6161e-02,  4.8017e-02,  1.1609e-01, -4.0465e-02, -1.0819e-01,\n",
      "         5.0446e-02,  9.5865e-02, -2.1391e-02, -7.7671e-02, -3.9044e-02,\n",
      "         1.3844e-01, -6.5074e-02, -1.4867e-01, -1.9490e-02, -1.7936e-02,\n",
      "        -5.1626e-03,  1.8408e-02, -6.9146e-02,  4.7459e-03, -4.2181e-02,\n",
      "        -7.8580e-02, -8.7683e-02,  1.3931e-01,  9.0899e-02,  2.8950e-02,\n",
      "        -5.3773e-02, -1.0698e-01,  1.0230e-02, -2.2217e-02,  1.0452e-01,\n",
      "        -2.0361e-01,  2.3256e-02,  7.0329e-02, -2.5763e-02,  1.7987e-01,\n",
      "         1.1091e-01,  1.3019e-01,  5.1247e-03,  1.7735e-01,  3.8684e-02,\n",
      "         8.6837e-02, -7.3407e-02, -1.6919e-01,  3.8702e-02,  8.1221e-03,\n",
      "        -2.5654e-02,  1.5659e-01, -2.9260e-03,  5.1470e-02, -3.8441e-02,\n",
      "        -4.2217e-02,  2.2510e-02, -8.9188e-02,  4.8149e-02, -8.4771e-02,\n",
      "         4.9795e-04,  1.0338e-01, -7.2772e-02, -6.4314e-02, -9.7352e-02,\n",
      "        -5.1070e-02, -5.8602e-02, -1.5343e-01, -2.8229e-02,  4.3786e-02,\n",
      "        -1.6760e-02,  4.3468e-02,  6.1846e-03,  1.0472e-01, -2.2352e-01,\n",
      "        -1.1572e-01,  6.7032e-02, -7.1199e-02, -8.9972e-02, -3.8627e-03,\n",
      "        -1.0734e-02,  1.5409e-02, -1.3571e-01, -1.4936e-01,  7.1428e-02,\n",
      "        -3.5436e-02,  3.4851e-02, -3.9899e-03, -1.3764e-03, -5.0805e-02,\n",
      "        -7.8808e-02, -1.6819e-01,  7.7438e-02, -4.7113e-02, -5.4881e-02,\n",
      "         6.3641e-02, -3.0046e-02,  1.0795e-02, -2.7979e-02, -2.7866e-02,\n",
      "         5.8352e-02,  1.1758e-01, -8.3546e-02, -2.7689e-02, -9.0345e-02,\n",
      "        -1.8160e-01,  1.1237e-01, -2.0741e-02,  8.0585e-02, -4.7589e-02,\n",
      "        -6.6311e-02,  2.0391e-02, -1.1035e-01, -6.6758e-02,  2.2329e-02,\n",
      "        -7.5694e-02,  2.8022e-02, -6.4829e-02, -9.0639e-02, -9.2730e-02,\n",
      "        -1.0422e-01,  4.8670e-03,  7.1762e-02, -6.5062e-03,  6.3636e-02,\n",
      "         3.2532e-03, -1.1631e-01, -1.3105e-01, -1.2366e-02,  5.0156e-02,\n",
      "         5.4343e-02, -2.1215e-02,  6.7260e-02, -1.3586e-01,  2.7010e-02,\n",
      "        -8.6845e-02, -9.0624e-02,  8.3127e-02,  3.8983e-02,  4.5054e-02,\n",
      "        -1.5408e-01,  9.0453e-03,  1.2971e-02, -5.0808e-02,  3.5881e-02,\n",
      "         3.9825e-02, -3.7674e-02, -4.8837e-02, -1.3166e-01, -1.8676e-02,\n",
      "         5.9037e-02, -1.2694e-01,  2.0501e-02,  1.6616e-01, -3.5600e-02,\n",
      "         3.9715e-02,  1.2555e-01, -1.2326e-01, -1.0526e-01, -1.5849e-02,\n",
      "         1.0186e-01, -9.6679e-02, -1.0716e-02, -4.6208e-02,  9.3579e-02,\n",
      "        -2.8423e-02, -4.4852e-03,  1.6591e-01,  9.7545e-02, -4.4024e-03,\n",
      "        -3.0272e-02,  5.0326e-02, -6.5839e-02,  1.6873e-01, -1.0768e-01,\n",
      "         8.5528e-02,  1.6620e-02, -5.8651e-02,  3.7556e-02, -7.5479e-02,\n",
      "        -1.2982e-01, -1.7412e-02,  3.1357e-02,  1.3156e-01,  1.2501e-02,\n",
      "         6.4578e-05,  1.1438e-02, -3.2406e-02, -8.3234e-03, -4.0322e-02,\n",
      "         3.8501e-02,  1.4806e-01,  5.0519e-02, -5.9282e-02, -1.1720e-01,\n",
      "        -2.2978e-02,  2.3423e-02, -1.9841e-02, -1.2046e-01, -7.9394e-03,\n",
      "         2.6962e-02,  3.2584e-02,  3.4176e-02, -1.7455e-02,  5.4376e-02,\n",
      "        -6.2332e-02,  8.9045e-02, -1.2350e-01, -4.2116e-02, -4.1007e-02,\n",
      "        -1.0059e-01, -9.9187e-02, -3.1925e-02, -1.6104e-01,  3.8237e-02,\n",
      "         6.7808e-02, -2.1925e-02, -3.5237e-02,  1.6929e-03,  5.3257e-02,\n",
      "        -4.9049e-02, -1.3038e-02, -8.5554e-02, -6.6013e-02,  8.7831e-02,\n",
      "         1.0763e-02,  4.7769e-02, -5.3345e-02, -1.4700e-01, -8.2749e-02,\n",
      "         7.5559e-03,  7.3252e-02, -8.5247e-02, -1.2353e-02, -3.6422e-02,\n",
      "        -3.1804e-02,  1.3748e-01,  6.8776e-02,  6.9719e-02, -6.1513e-02,\n",
      "         1.3168e-01, -5.5860e-02,  1.2063e-01,  3.0924e-02, -1.8618e-01,\n",
      "         9.1942e-02, -5.7907e-02,  4.5920e-02,  1.4161e-02,  5.4645e-03,\n",
      "        -1.0145e-01,  1.4069e-02,  7.0383e-02,  3.3249e-02, -7.1238e-02,\n",
      "        -9.5610e-02, -8.9521e-03,  7.9892e-02, -5.9931e-02,  1.3094e-02,\n",
      "         8.4902e-02, -8.8815e-02, -3.8338e-02, -2.3316e-03, -5.8273e-02,\n",
      "         3.5278e-03,  1.8488e-02,  6.8640e-02, -1.5735e-01, -9.4957e-02,\n",
      "        -3.4438e-02,  2.0991e-01, -1.9748e-02,  1.4705e-01, -3.4270e-02,\n",
      "        -6.3943e-02,  5.7946e-02,  4.6824e-02, -4.7066e-02,  1.8138e-02,\n",
      "        -1.0371e-01, -2.7183e-02, -3.4263e-02, -4.1248e-02, -1.1747e-01,\n",
      "        -2.6096e-04,  6.9966e-02,  2.9752e-02,  5.5570e-02, -3.4749e-02,\n",
      "        -4.9083e-02,  1.5480e-01, -2.5728e-02, -1.4297e-01, -1.8912e-02,\n",
      "        -2.5677e-02,  4.3914e-02, -5.8983e-02, -7.6274e-02, -1.8738e-02,\n",
      "        -4.7779e-02,  5.3665e-02, -9.5376e-03,  2.2442e-02,  8.5170e-02,\n",
      "         1.3105e-01, -3.2606e-02, -1.0897e-01,  8.2333e-02, -1.7728e-01,\n",
      "        -2.2134e-02,  1.6195e-01, -9.9854e-03,  7.9816e-02, -4.8753e-03])\n",
      "mpn.W_i.weight tensor([[-0.1408,  0.0116, -0.2053,  ...,  0.2381, -0.0154, -0.0421],\n",
      "        [-0.3682,  0.1187,  0.2028,  ..., -0.1483,  0.0113,  0.0239],\n",
      "        [-0.0705, -0.0354, -0.1530,  ...,  0.0388, -0.0766, -0.0558],\n",
      "        ...,\n",
      "        [ 0.0694, -0.5387,  0.0529,  ..., -0.0906, -0.0105,  0.1408],\n",
      "        [-0.0478, -0.0072, -0.0364,  ..., -0.2287,  0.0091, -0.0168],\n",
      "        [-0.1602,  0.0714,  0.2547,  ...,  0.1519,  0.1099, -0.0915]])\n",
      "mpn.W_h.weight tensor([[ 0.2999, -0.0738,  0.0767,  ...,  0.0246,  0.0260,  0.1108],\n",
      "        [-0.0381, -0.3316, -0.2135,  ...,  0.0364,  0.0386, -0.2138],\n",
      "        [ 0.0335, -0.0062, -0.2131,  ...,  0.1398, -0.1080, -0.1794],\n",
      "        ...,\n",
      "        [ 0.0517,  0.1110,  0.1762,  ..., -0.2135, -0.1413, -0.0790],\n",
      "        [-0.1140,  0.1367, -0.0972,  ..., -0.0510, -0.0074,  0.0840],\n",
      "        [-0.1956,  0.0779, -0.5473,  ..., -0.2423, -0.0690,  0.1945]])\n",
      "mpn.W_o.weight tensor([[ 0.0110,  0.1306, -0.2996,  ..., -0.3920,  0.0236, -0.0939],\n",
      "        [ 0.0670, -0.5404,  0.1546,  ..., -0.2579,  0.1009, -0.0461],\n",
      "        [-0.1829,  0.0440,  0.1627,  ...,  0.0509, -0.1251, -0.1174],\n",
      "        ...,\n",
      "        [ 0.0543,  0.1307,  0.0483,  ..., -0.3569,  0.1105,  0.0067],\n",
      "        [ 0.0834, -0.0528, -0.0529,  ..., -0.0909, -0.1027,  0.0950],\n",
      "        [-0.0318, -0.0256,  0.1212,  ...,  0.1489,  0.0200,  0.0284]])\n",
      "mpn.W_o.bias tensor([ 0.0306,  0.0169, -0.0311,  0.0118, -0.0492,  0.0173, -0.0953,  0.0334,\n",
      "        -0.0272, -0.0425,  0.0522, -0.0560, -0.0366, -0.0625,  0.0080, -0.1244,\n",
      "        -0.0043, -0.0476, -0.1038, -0.0790, -0.1285, -0.0702, -0.0232, -0.0251,\n",
      "        -0.0725, -0.0441, -0.0689, -0.0476, -0.0197, -0.0440, -0.0016, -0.0875,\n",
      "         0.0017, -0.0369,  0.0368, -0.0376,  0.0017, -0.0912, -0.0114, -0.0388,\n",
      "        -0.0270, -0.0464,  0.0258,  0.0725,  0.0012, -0.0923,  0.0187, -0.0124,\n",
      "         0.0062, -0.0060,  0.0067, -0.0527, -0.0245, -0.0290, -0.0239, -0.0989,\n",
      "        -0.0302, -0.0744, -0.0225, -0.0872, -0.0409,  0.0174, -0.0701, -0.0530,\n",
      "        -0.0900, -0.2223, -0.0605,  0.0501, -0.0507, -0.0384, -0.0407, -0.0610,\n",
      "        -0.0188, -0.0164,  0.0284, -0.0218, -0.1055, -0.0121, -0.0502, -0.0148,\n",
      "        -0.0225, -0.0424,  0.0436, -0.0356,  0.0125, -0.1154, -0.0673, -0.0627,\n",
      "        -0.0802, -0.0951,  0.0456, -0.0629, -0.0649, -0.0064,  0.0044, -0.1329,\n",
      "        -0.0004,  0.0351, -0.0660, -0.0730,  0.0053, -0.0047, -0.0248,  0.0076,\n",
      "        -0.0496, -0.0337, -0.0804, -0.0312,  0.0343, -0.1383,  0.0582,  0.0218,\n",
      "        -0.0550, -0.0204, -0.0416, -0.0276, -0.0272, -0.0595, -0.0457, -0.0655,\n",
      "        -0.0028, -0.0473, -0.0819, -0.0181, -0.0347, -0.0832, -0.0047, -0.0737,\n",
      "        -0.0101,  0.0464, -0.0230, -0.0563, -0.0579, -0.0142, -0.0405, -0.0175,\n",
      "        -0.0390, -0.1181, -0.0539, -0.2084, -0.0191, -0.0867, -0.0973, -0.0689,\n",
      "        -0.0341,  0.0279, -0.0528, -0.0086,  0.0066, -0.0596, -0.0783, -0.0720,\n",
      "        -0.0290, -0.1218, -0.0703, -0.0666, -0.0662, -0.0583, -0.0300, -0.0366,\n",
      "        -0.0644, -0.0249, -0.0427, -0.0710, -0.0705, -0.0035, -0.0399,  0.0335,\n",
      "         0.0162,  0.0480,  0.0007, -0.0510, -0.0359,  0.0188,  0.0443, -0.0678,\n",
      "        -0.0204, -0.0345, -0.0005,  0.0108, -0.0094, -0.0658, -0.0834, -0.0016,\n",
      "        -0.0479,  0.0366,  0.0411, -0.0551,  0.0131,  0.0023, -0.0093, -0.0120,\n",
      "        -0.0459, -0.0667, -0.0347, -0.0241, -0.0392, -0.1025, -0.0024, -0.0229,\n",
      "        -0.0498,  0.0082,  0.0035, -0.0255, -0.0767, -0.0090, -0.0180, -0.0573,\n",
      "        -0.0177, -0.0520, -0.0296, -0.0449, -0.0980,  0.0306, -0.0257, -0.0790,\n",
      "        -0.0926, -0.0348, -0.1044, -0.0655, -0.0432, -0.0161, -0.0263, -0.0423,\n",
      "        -0.0410,  0.0037, -0.0876, -0.0381, -0.1037, -0.0954, -0.0332,  0.0256,\n",
      "         0.0298,  0.0145, -0.0164,  0.0105, -0.0762, -0.0831, -0.0041, -0.0871,\n",
      "        -0.0770, -0.0266, -0.0500,  0.0205,  0.0211, -0.0341, -0.0223, -0.0886,\n",
      "        -0.1174, -0.0519, -0.0615, -0.0210, -0.1131, -0.0505, -0.0703, -0.0984,\n",
      "        -0.0510, -0.0303,  0.0292, -0.0060, -0.0419, -0.0679,  0.0205, -0.1140,\n",
      "        -0.0042, -0.0564, -0.0134, -0.0487, -0.0447, -0.0699, -0.1627,  0.0138,\n",
      "         0.0159,  0.0110,  0.0128,  0.0219, -0.0523, -0.0096, -0.0151, -0.0242,\n",
      "        -0.0711,  0.0181,  0.0140, -0.0137,  0.0415, -0.0467, -0.0469, -0.0019,\n",
      "         0.0458,  0.0271, -0.0640,  0.0087, -0.0164,  0.0370, -0.0602, -0.0463,\n",
      "         0.0575, -0.1036, -0.1132, -0.0266, -0.0679, -0.0967, -0.1371,  0.0018,\n",
      "        -0.0204, -0.0395,  0.0125, -0.0793, -0.0557, -0.1040,  0.0483, -0.0271,\n",
      "        -0.0395, -0.0475, -0.0821, -0.0938, -0.0386, -0.0238, -0.0342, -0.0402,\n",
      "         0.0316,  0.0400,  0.0158, -0.0097, -0.0514, -0.1028, -0.0601, -0.0163,\n",
      "         0.0070, -0.0337, -0.0421, -0.0798, -0.0183,  0.0331, -0.1007, -0.0732,\n",
      "        -0.0045, -0.0690,  0.0012, -0.0011, -0.0331, -0.0782,  0.0223, -0.0137,\n",
      "        -0.0045, -0.0510, -0.0663, -0.0705, -0.0094, -0.0641, -0.0087, -0.0062,\n",
      "        -0.0315, -0.0355, -0.1104, -0.0751, -0.0187,  0.0014, -0.0517, -0.0029,\n",
      "         0.0141, -0.0872,  0.0755, -0.0296,  0.0629, -0.1056, -0.0591, -0.0039,\n",
      "        -0.0968, -0.0814, -0.0805, -0.1048, -0.0951, -0.0422, -0.1094, -0.0385,\n",
      "        -0.0272, -0.0227, -0.0566, -0.0110, -0.1241, -0.0173, -0.0678, -0.1003,\n",
      "        -0.1465, -0.1234, -0.0088, -0.0517, -0.1041,  0.0117, -0.0336,  0.0233,\n",
      "        -0.0747, -0.0238, -0.0149, -0.0263, -0.0789, -0.0361, -0.0086,  0.0153,\n",
      "        -0.0636, -0.0432, -0.0315,  0.0351,  0.0240, -0.0415, -0.0036, -0.0133,\n",
      "        -0.0322, -0.2164, -0.0394,  0.0266, -0.0173, -0.0179, -0.0289, -0.0430,\n",
      "        -0.0013,  0.0055, -0.0145,  0.0359,  0.0132, -0.0851,  0.0128, -0.0225,\n",
      "        -0.0851,  0.0466,  0.0319, -0.0369, -0.0660, -0.0505, -0.0557,  0.0030,\n",
      "         0.0046,  0.0671, -0.0734, -0.0253, -0.0577, -0.0603, -0.0199,  0.0066,\n",
      "        -0.0141, -0.0114, -0.0660, -0.0525, -0.1804, -0.0432, -0.0243, -0.0563,\n",
      "         0.0046, -0.0260])\n",
      "T_mean.weight tensor([[ 0.1098, -0.0040,  0.0612,  ...,  0.0667, -0.0276,  0.0482],\n",
      "        [-0.3186,  0.0239, -0.0576,  ...,  0.1034, -0.0456,  0.1216],\n",
      "        [ 0.3669, -0.1705, -0.0551,  ..., -0.0468, -0.1341,  0.1362],\n",
      "        ...,\n",
      "        [-0.0674,  0.1682,  0.2772,  ...,  0.1152,  0.0067, -0.0249],\n",
      "        [-0.3291,  0.0303,  0.0291,  ...,  0.0136, -0.0654, -0.0810],\n",
      "        [-0.0986,  0.0578, -0.0170,  ..., -0.0604, -0.0049, -0.1958]])\n",
      "T_mean.bias tensor([-0.0720, -0.0099,  0.0569,  0.1031, -0.0147, -0.0331,  0.0500,  0.0883,\n",
      "         0.0395, -0.0898, -0.0089, -0.0547,  0.0396, -0.0502, -0.0595,  0.0227,\n",
      "        -0.1039,  0.1130,  0.0130, -0.0018,  0.0684, -0.1394, -0.0162,  0.0078,\n",
      "        -0.0738, -0.0177,  0.0426,  0.0395])\n",
      "T_var.weight tensor([[-0.0010, -0.0828,  0.0483,  ..., -0.0211,  0.1110, -0.0958],\n",
      "        [ 0.3346, -0.0488,  0.0484,  ...,  0.0063, -0.0803,  0.0339],\n",
      "        [-0.1311, -0.0643, -0.0481,  ...,  0.0967,  0.1810, -0.0791],\n",
      "        ...,\n",
      "        [ 0.0601,  0.0208, -0.0351,  ..., -0.0330,  0.0800,  0.0722],\n",
      "        [ 0.1391, -0.0895,  0.0046,  ...,  0.0385,  0.0048,  0.0808],\n",
      "        [-0.0651, -0.0145,  0.0648,  ...,  0.1260, -0.1752, -0.0567]])\n",
      "T_var.bias tensor([-0.0441, -0.0728, -0.0512,  0.0758, -0.0763,  0.0137, -0.1198,  0.0703,\n",
      "        -0.0790,  0.1000,  0.0582, -0.0382,  0.0524, -0.0797, -0.0990, -0.0882,\n",
      "         0.0623, -0.1038, -0.0598,  0.1289, -0.0584, -0.0485,  0.0501,  0.0430,\n",
      "         0.0768,  0.0630, -0.0500, -0.0909])\n",
      "G_mean.weight tensor([[-0.0557,  0.1280, -0.0015,  ...,  0.0578,  0.0991, -0.2037],\n",
      "        [ 0.0660, -0.2612,  0.1959,  ...,  0.0097,  0.1692, -0.0145],\n",
      "        [-0.1203,  0.0115,  0.1970,  ..., -0.0727,  0.1105,  0.0473],\n",
      "        ...,\n",
      "        [-0.1914,  0.2078, -0.1512,  ...,  0.1153,  0.0080, -0.0011],\n",
      "        [ 0.1169,  0.1587,  0.0239,  ...,  0.0942, -0.0528, -0.0406],\n",
      "        [-0.1011, -0.1089, -0.0350,  ...,  0.0317,  0.0181, -0.0484]])\n",
      "G_mean.bias tensor([ 0.0464, -0.0436,  0.0857,  0.0175,  0.0703,  0.0063,  0.0459,  0.0221,\n",
      "         0.0357, -0.0713, -0.0553, -0.0744,  0.0492,  0.1520, -0.0717,  0.0342,\n",
      "         0.0680, -0.1086,  0.0377,  0.0223, -0.0360,  0.0428,  0.0132, -0.0642,\n",
      "        -0.0761,  0.0232, -0.0227,  0.0625])\n",
      "G_var.weight tensor([[ 0.5010,  0.0648, -0.1036,  ...,  0.1036,  0.1769,  0.1902],\n",
      "        [-0.6629, -0.1006, -0.0161,  ..., -0.2135, -0.2141, -0.0646],\n",
      "        [ 0.4921,  0.1824,  0.2322,  ...,  0.2878,  0.3563,  0.1567],\n",
      "        ...,\n",
      "        [-0.4589, -0.0729, -0.1435,  ..., -0.0396, -0.3327, -0.0898],\n",
      "        [ 0.3526,  0.1957,  0.0109,  ...,  0.2348,  0.2228, -0.0468],\n",
      "        [-0.5380, -0.3648, -0.1544,  ..., -0.1638, -0.0368, -0.1170]])\n",
      "G_var.bias tensor([ 0.4655, -0.5578,  0.5657,  0.4965, -0.3317,  0.6462,  0.5776,  0.5259,\n",
      "        -0.5419,  0.5308,  0.5359, -0.5042,  0.5566,  0.5319,  0.5368,  0.5904,\n",
      "        -0.5169,  0.6076, -0.6032, -0.5292, -0.4907,  0.4616, -0.5174, -0.4661,\n",
      "        -0.5509, -0.5612,  0.5845, -0.5638])\n",
      "ffn.1.weight tensor([[-0.0906, -0.3133,  0.1303,  ..., -0.1678, -0.2319, -0.2223],\n",
      "        [ 0.0267,  0.0584,  0.0042,  ..., -0.0318,  0.2302, -0.1829],\n",
      "        [ 0.0468, -0.0302, -0.2179,  ...,  0.0975,  0.0639,  0.0909],\n",
      "        ...,\n",
      "        [ 0.0759, -0.1012, -0.1698,  ..., -0.0617, -0.1279,  0.2008],\n",
      "        [ 0.0345, -0.0945,  0.0157,  ...,  0.0509,  0.0936,  0.1389],\n",
      "        [-0.1152,  0.0739, -0.0028,  ...,  0.1512, -0.1347, -0.1339]])\n",
      "ffn.1.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.])\n",
      "ffn.4.weight tensor([[ 0.1104, -0.0925,  0.0434,  ...,  0.1937, -0.0719,  0.0526],\n",
      "        [ 0.0602, -0.1469, -0.1582,  ..., -0.0870,  0.0824,  0.2012],\n",
      "        [-0.1320,  0.1234,  0.1687,  ..., -0.1030, -0.0078, -0.0378],\n",
      "        ...,\n",
      "        [ 0.1929,  0.1886,  0.0030,  ..., -0.0433, -0.1526,  0.0185],\n",
      "        [ 0.1433,  0.0100,  0.3022,  ...,  0.0065, -0.0185, -0.0090],\n",
      "        [-0.0837,  0.1812, -0.0791,  ...,  0.0365, -0.2549,  0.2905]])\n",
      "ffn.4.bias tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
      "        0., 0.])\n",
      "ffn.7.weight tensor([[-0.1975,  0.1618, -0.2512, -0.2361,  0.0810, -0.0652, -0.2099,  0.3093,\n",
      "          0.3060,  0.2265,  0.1251, -0.1239,  0.1891,  0.1764,  0.3673,  0.1286,\n",
      "          0.2934, -0.3947,  0.0574, -0.5267,  0.2699, -0.2439, -0.2550,  0.2081,\n",
      "          0.1409, -0.1278, -0.0208,  0.2078,  0.1396, -0.0285,  0.3486, -0.1499,\n",
      "          0.2901,  0.0223, -0.0744, -0.0607,  0.0075,  0.1774, -0.2489, -0.3042,\n",
      "          0.1305,  0.2034,  0.1850, -0.0236,  0.0462, -0.1273,  0.1803, -0.1508,\n",
      "          0.0682, -0.0702]])\n",
      "ffn.7.bias tensor([0.])\n"
     ]
    }
   ],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print name, param.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model #Params: 2474K\n"
     ]
    }
   ],
   "source": [
    "model = model.cuda()\n",
    "print \"Model #Params: %dK\" % (sum([x.nelement() for x in model.parameters()]) / 1000,)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "scheduler = lr_scheduler.ExponentialLR(optimizer, 0.9)\n",
    "scheduler.step()\n",
    "\n",
    "dataset = MoleculeDataset(os.path.join(DATASET_PATH, 'logp_wo_averaging_validation.csv'), RAW_PATH)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Merge train and val datasets for CrossVal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_PATH = '../../../data/3_final_data/split_data'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = pd.read_csv(os.path.join(DATASET_PATH, 'logp_wo_averaging_train.csv'), index_col = 0)\n",
    "val_data = pd.read_csv(os.path.join(DATASET_PATH, 'logp_wo_averaging_validation.csv'), index_col = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>smiles</th>\n",
       "      <th>logP</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>CCN(CC)CCSC(=NO)c1noc(-c2ccccc2)n1</td>\n",
       "      <td>2.13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>COC(=O)NC1=NC(c2ccccc2Br)CN1</td>\n",
       "      <td>1.85</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                               smiles  logP\n",
       "1  CCN(CC)CCSC(=NO)c1noc(-c2ccccc2)n1  2.13\n",
       "3        COC(=O)NC1=NC(c2ccccc2Br)CN1  1.85"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data.iloc[[1,3]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "cv_data = pd.concat([train_data, val_data], axis = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_OUTPUT_PATH = \"../../../data/raw/baselines/jtree/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "cv_data.to_csv(os.path.join(DATASET_OUTPUT_PATH, 'logp_wo_averaging_train_val.csv'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train function draft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_EPOCH = 10\n",
    "PRINT_ITER = 20\n",
    "criterion = nn.MSELoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "tensor(21.4522, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(33.3940, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(4.6014, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "1\n",
      "tensor(61.3403, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(142.5458, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(50.6553, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "2\n",
      "tensor(22.7041, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(29.3614, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(42.1143, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "3\n",
      "tensor(30.2984, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(33.3286, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(56.2206, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "4\n",
      "tensor(43.3054, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(23.1365, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(23.1205, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "5\n",
      "tensor(74.3726, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(23.8956, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(29.0056, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "6\n",
      "tensor(38.2152, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(17.2063, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(73.9667, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "7\n",
      "tensor(84.9256, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(48.4532, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(52.9833, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "8\n",
      "tensor(6.1152, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(175.1078, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(21.1776, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "9\n",
      "tensor(12.5332, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(22.7721, device='cuda:0', grad_fn=<MeanBackward1>)\n",
      "tensor(11.1780, device='cuda:0', grad_fn=<MeanBackward1>)\n"
     ]
    }
   ],
   "source": [
    "for epoch in xrange(MAX_EPOCH):\n",
    "    dataloader = DataLoader(dataset, batch_size=3, shuffle=True, collate_fn=lambda x: x, num_workers=1, drop_last=True)\n",
    "\n",
    "    print(epoch)\n",
    "\n",
    "    for it,  batch in enumerate(dataloader):\n",
    "        X = []\n",
    "        y = []\n",
    "        for elem in batch:\n",
    "            X.append(elem[0])\n",
    "            y.append(elem[1])\n",
    "        y = torch.Tensor(y).cuda()\n",
    "# #         batch, target = batch\n",
    "# #         for mol_tree in batch:\n",
    "# #             for node in mol_tree.nodes:\n",
    "# #                 if node.label not in node.cands:\n",
    "# #                     node.cands.append(node.label)\n",
    "# #                     node.cand_mols.append(node.label_mol)\n",
    "\n",
    "# #         try:\n",
    "        model.zero_grad()\n",
    "        pred = model(X)\n",
    "        loss = criterion(y, pred)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        print(loss)\n",
    "#         print('sdfsd')\n",
    "#         except Exception as e:\n",
    "#             with open('broken_smiles.txt', 'a') as f:\n",
    "#                 smiles = [elem.smiles for elem in batch]\n",
    "#                 f.write('\\n'.join(smiles))\n",
    "#                 f.write('\\n')\n",
    "#                 f.write(e.args[0])\n",
    "#                 f.write('\\n\\n\\n\\n')\n",
    "#             print(e)\n",
    "#             continue\n",
    "\n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'CC(=O)OCCC(C)C': <mol_tree.MolTree at 0x7efbabb36d90>,\n",
       " 'CC(=O)c1ccc2cc(C(C)C(=O)OCC(=O)N(C)C)ccc2c1': <mol_tree.MolTree at 0x7efbaaa1d750>,\n",
       " 'CC(C)(C)S(=O)(=O)CC(Cc1ccccc1)C(=O)NC(Cc1c[nH]cn1)C(=O)NC(CC1CCCCC1)C(O)C(O)C1CC1': <mol_tree.MolTree at 0x7efbaaa7b090>,\n",
       " 'CN=C(NC#N)NCCSCc1nccs1': <mol_tree.MolTree at 0x7efbabb68750>,\n",
       " 'Clc1ccc(-c2nc3cccnc3[nH]2)cc1': <mol_tree.MolTree at 0x7efbaaa7bfd0>,\n",
       " 'N#CN=[N+]([O-])c1ccc(Br)cc1': <mol_tree.MolTree at 0x7efbabb685d0>,\n",
       " 'N=c1nc(-c2ccccc2Br)[nH]c(=N)[nH]1': <mol_tree.MolTree at 0x7efbafced490>,\n",
       " 'O=C1CCc2ccccc2N1': <mol_tree.MolTree at 0x7efbaaa7bc90>,\n",
       " 'O=C1c2ccccc2C(=O)N1SC(Cl)(Cl)Cl': <mol_tree.MolTree at 0x7efbabb85690>,\n",
       " 'O=c1ccc(=O)[nH][nH]1': <mol_tree.MolTree at 0x7efbaaa7b050>}"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SMILES_TO_MOLTREE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('N#CN=[N+]([O-])c1ccc(Br)cc1', {})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(<mol_tree.MolTree at 0x7efbabb61f10>, 2.5)"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('O=C1c2ccccc2C(=O)N1SC(Cl)(Cl)Cl', {'N#CN=[N+]([O-])c1ccc(Br)cc1': <mol_tree.MolTree object at 0x7efbabb61f10>})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(<mol_tree.MolTree at 0x7efbaa980350>, 2.85)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('CC(C)(C)S(=O)(=O)CC(Cc1ccccc1)C(=O)NC(Cc1c[nH]cn1)C(=O)NC(CC1CCCCC1)C(O)C(O)C1CC1', {'N#CN=[N+]([O-])c1ccc(Br)cc1': <mol_tree.MolTree object at 0x7efbabb61f10>, 'O=C1c2ccccc2C(=O)N1SC(Cl)(Cl)Cl': <mol_tree.MolTree object at 0x7efbaa980350>})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(<mol_tree.MolTree at 0x7efbaa980510>, 2.75)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('N#CN=[N+]([O-])c1ccc(Br)cc1', {'N#CN=[N+]([O-])c1ccc(Br)cc1': <mol_tree.MolTree object at 0x7efbabb61f10>, 'CC(C)(C)S(=O)(=O)CC(Cc1ccccc1)C(=O)NC(Cc1c[nH]cn1)C(=O)NC(CC1CCCCC1)C(O)C(O)C1CC1': <mol_tree.MolTree object at 0x7efbaa980510>, 'O=C1c2ccccc2C(=O)N1SC(Cl)(Cl)Cl': <mol_tree.MolTree object at 0x7efbaa980350>})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(<mol_tree.MolTree at 0x7efbabb61f10>, 2.5)"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:jtree]",
   "language": "python",
   "name": "conda-env-jtree-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}