Upload BiLSTM

blstm_config.py

@@ -5,7 +5,7 @@ import torch
 class BiLSTMConfig(PretrainedConfig):
     def __init__(self, vocab_size=23626, embed_dim=100,
                  num_layers=1, hidden_dim=256, dropout=0.33,
-                 output_dim=128, predict_output=10, **kwargs):
+                 output_dim=128, predict_output=10, device="cuda:0", **kwargs):
         
         super().__init__(**kwargs)
         self.vocab_size = vocab_size
@@ -14,4 +14,5 @@ class BiLSTMConfig(PretrainedConfig):
         self.hidden_dim = hidden_dim
         self.dropout = dropout
         self.output_dim = output_dim
-        self.predict_output = predict_output
+        self.predict_output = predict_output
+        self.device = device

blstm_model.py

@@ -19,7 +19,7 @@ class BiLSTM(PreTrainedModel):
         self.dropout = nn.Dropout(config.dropout)
         self.elu = nn.ELU()
         self.fc = nn.Linear(config.output_dim, config.predict_output)
-        # self.device_ = config.device
+        self.device_ = config.device
     
     def forward(self, input):   # input is a list of indices, shape batch_size, seq_len
         x = self.embed_layer(input)                     # batch_size, seq_len, 100  (This is only when batch_first=True!!!!)
@@ -37,6 +37,6 @@ class BiLSTM(PreTrainedModel):
         return out, hidden
     
     def init_hidden(self, batch_size):
-        hidden = torch.zeros(2, batch_size, self.hidden_dim//2)
-        cell = torch.zeros(2, batch_size, self.hidden_dim//2)
+        hidden = torch.zeros(2, batch_size, self.hidden_dim//2, device=self.device_)
+        cell = torch.zeros(2, batch_size, self.hidden_dim//2, device=self.device_)
         return hidden, cell

config.json

@@ -6,6 +6,7 @@
     "AutoConfig": "blstm_config.BiLSTMConfig",
     "AutoModel": "blstm_model.BiLSTM"
   },
+  "device": "cuda:0",
   "dropout": 0.33,
   "embed_dim": 100,
   "hidden_dim": 256,