My Google Colab keeps crashing because of something in this code. Not sure what it is
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. 

This code is followed by generate function, encoder decoder class, etc. It is when I get to train that it crashes (for "unknown reason")

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 



The solution was to add cuda() after defining each tensor. 

E.g. 

a = a.cuda()


Note that you must assign it instead of doing just a.cuda()
