由于这段代码中的某些内容，我的Google Colab一直崩溃。不知道是什么

我的Google colab一直在火车上崩溃，即使RAM和磁盘很多。我认为问题出在代码中的某个地方，但我不知道它是什么。我在做LSTM。我将不胜感激。我正在使用PyTorch。

此代码后跟生成函数，编码器，解码器类等。这是在我训练时崩溃的原因（由于“未知原因”）

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