'The error of several nn.GRU layers in Pytorch model

Recently, I made a GRU model with Pytorch. When the model has one nn.GRU layer, it runs well. But when there are more than one GRU layer, the model would report error like ' CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle)' at the start of training. The detailed model architecture is attached at the end. My CUDA vesion is 11.0 and Pytorch vesion is 1.7. Does it mean I can't use GRU model with more than one nn.GRU layer?

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


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source