Use the dataset class to initialize model parameters

1 files changed, 3 insertions, 2 deletions

diff --git a/model.py b/model.py
index 9b08f47..cfa1deb 100644
--- a/model.py
+++ b/model.py
@@ -2,13 +2,14 @@ import torch.nn as nn
 
 
 class Model(nn.Module):
-    def __init__(self, n_vocab, seq_size, embedding_size, lstm_size, num_layers, dropout):
+    def __init__(self, dataset, embedding_size, lstm_size, num_layers, dropout):
         super(Model, self).__init__()
         
-        self.seq_size = seq_size
+        self.seq_size = dataset.seq_size
         self.lstm_size = lstm_size
         self.num_layers = num_layers
 
+        n_vocab = len(dataset.uniq_words)
         self.embedding = nn.Embedding(n_vocab, embedding_size)
         
         self.lstm = nn.LSTM(