Сбой PyTorch LSTM на GPU Colab (отлично работает на CPU)

Question

Здравствуйте, у меня есть следующий LSTM, который отлично работает на процессоре.

Однако, когда я пытаюсь запустить на Colab GPU, он вылетает даже без сообщения об ошибке.

   model = LSTMForecast(1, 1, batch_size=1).to('cuda')
   a = torch.rand(1, 1, 1).to('cuda')
   model(a)

Логи мне тоже ничего не говорят. Я действительно в растерянности.

Answer 1

Мне пришлось явно позвонить в CUDA. Как только я это сделал, это сработало.

def init_hidden(self, batch_size)->None:
    # This is what we'll initialise our hidden state
    self.hidden = (torch.zeros(self.num_layers, batch_size, self.hidden_dim).to('cuda'), torch.zeros(self.num_layers, batch_size, self.hidden_dim).to('cuda'))