Определить плоский слой в нейронной сети с помощью pytorch

Я пытаюсь определить плоский слой перед запуском полносвязного слоя. Поскольку мои входные данные представляют собой тензор формы (512, 2, 2), поэтому я хочу сгладить этот тензор перед слоями FC.

Раньше я получал эту ошибку:

empty(): argument 'size' must be tuple of ints, but found element of type Flatten at pos 2

import torch.nn as nn
class Network(nn.Module):
    def __init__(self):
        super(Network,self).__init__()
        self.flatten=nn.Flatten()
        self.fc1=nn.Linear(self.flatten,512)
        self.fc2=nn.Linear(512,256)
        self.fc3=nn.Linear(256,3)
 
        
    def forward(self,x):
        x=self.flatten(x) # Flatten layer
        x=torch.ReLU(self.fc1(x))  
        x=torch.ReLU(self.fc2(x))
        x=torch.softmax(self.fc3(x))
        return x