title: My Google Colab keeps crashing because of something in this code. Not sure what it is
body: <p>My google colab keeps crashing at train, even though RAM and disk are plenty. The problem is somewhere in this code, I think, but I don't know what it is. I am doing LSTM. I'd appreciate any help. I am using PyTorch. </p>
  
  <p>This code is followed by generate function, encoder decoder class, etc. It is when I get to train that it crashes (for "unknown reason")</p>
  
  <pre class="lang-py prettyprint-override"><code>class LSTMLM(torch.nn.Module):
    def __init__(self,
                vocab_size,
                embedding_size,
                hidden_size,
                num_layers=1,
                dropout=0.1):
      super().__init__()
      self.vocab_size = vocab_size
      self.embedding_size = embedding_size
      self.hidden_size = hidden_size
      self.num_layers = num_layers
      self.dropout = torch.nn.Dropout(dropout)
  
      self.embedding = torch.nn.Embedding(vocab_size, embedding_size)
  
      self.lstm = torch.nn.LSTM(input_size=hidden_size, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
  
      self.logistl = torch.nn.Linear(hidden_size, vocab_size)
  
      pass
  
    def forward(self, x, init_hidden_state=None):
      assert x.shape[0] == 1
  
      emb = self.embedding(x)
  
      emb = self.dropout(emb)
  
      if init_hidden_state is None:
        h0 = torch.zeros(self.num_layers, 1, self.hidden_size)
        c0 = torch.zeros(self.num_layers, 1, self.hidden_size)
      else:
        h0, c0 = init_hidden_state
      output, (hn, hc) = self.lstm(emb, (h0, c0))
  
      hidden_states = output
      final_hidden_state = hn
      final_cell_state = cn
      final_state = [final_hidden__state, final_cell_state]
  
      hidden_states = self.dropout(hidden_states) 
  
      output_dist = self.logistl(hidden_states)
  
      return output_dist, hidden_states, final_state 
  
  </code></pre>