text: #zz import torch
  #zz from torch import nn
  import tensorflow as tf
  import tensorflow.nn as nn 
  from gentrl.tokenizer import encode, get_vocab_size
  
  
  class RNNEncoder(nn.Module):
      def __init__(self, hidden_size=256, num_layers=2, latent_size=50,
                   bidirectional=False):
          super(RNNEncoder, self).__init__()
  
          self.embs = nn.Embedding(get_vocab_size(), hidden_size)
          self.rnn = nn.GRU(input_size=hidden_size,
                            hidden_size=hidden_size,
                            num_layers=num_layers,
                            bidirectional=bidirectional)
  
          self.final_mlp = nn.Sequential(
              nn.Linear(hidden_size, hidden_size), nn.LeakyReLU(),
              nn.Linear(hidden_size, 2 * latent_size))
  
      def encode(self, sm_list):
          """
          Maps smiles onto a latent space
          """
  
          tokens, lens = encode(sm_list)
          to_feed = tokens.transpose(1, 0).to(self.embs.weight.device)
  
          outputs = self.rnn(self.embs(to_feed))[0]
          outputs = outputs[lens, torch.arange(len(lens))]
  
          return self.final_mlp(outputs)
  