Wrong tensor size in GRU

I am trying to create a GRU but I am running into a tensor size problem. The shape of the input tensor is (16, 3, 256, 256). But when I try to run the GRU I get an error on self.GRU(out) that says RuntimeError : input must have 3 dimensions, got 4 and I believe this is because my out shape is (16, 3, 127, 127) but I don’t know I how would go about changing this to the right shape.

class Net(nn.Module):

def __init__(self):
    super(Net, self).__init__()
    self.pool = nn.MaxPool2d(4, 2)
    self.gru = nn.GRU(input_size = 128, hidden_size=hidden_size, bidirectional=True)

def forward(self, x):
    out = self.pool(x)
    out = torch.squeeze(out, -1)
    out = self.gru(out)

From the docs:

input of shape (seq_len, batch, input_size): tensor containing the features of the input sequence. The input can also be a packed variable length sequence. See torch.nn.utils.rnn.pack_padded_sequence() for details.

So indeed a 3-dimensional input tensor is expected in the mentioned shape.