Ошибка нескольких слоев nn.GRU в модели Pytorch

Question

Недавно я сделал модель ГРУ с помощью Pytorch. Когда модель имеет один слой nn.GRU, она работает хорошо. Но когда имеется более одного слоя GRU, модель будет сообщать об ошибке, например «ошибка CUDA: CUBLAS_STATUS_ALLOC_FAILED при вызове». ' в начале тренировки. Подробная архитектура модели прилагается в конце. Моя версия CUDA — 11.0, а версия Pytorch — 1.7. Означает ли это, что я не могу использовать модель GRU более чем с одним слоем nn.GRU?

      class GRU(nn.Module):
    def __init__(self): 
        super(GRU,self).__init__()
        self.gru1 = nn.GRU(3,50,1,batch_first=True)
        self.gru2 = nn.GRU(50,10,1,batch_first=True)
        self.fc1 = nn.Linear(10,2)

    def forward(self,x):
        x, self.hidsta1 = self.gru1(x)
        x, self.hidsta2 = self.gru2(x)
        s,b,h = x.shape
        x = x.reshape(s*b, h)
        x = self.fc1(x)
        x = x.reshape(s,b,-1)
        return x