Given input size: (512x1x1). Calculated output size: (512x0x0). Output size is too small

Есть вот такой алгоритм, получаю ошибку Given input size: (512x1x1). Calculated output size: (512x0x0). Output size is too small. Понимаю, что проблема в строке AvgPool. Не совпадают размеры. Не понимаю, как это исправить. Подскажите, пожалуйста.

class AntispoofModel(nn.Module):
    def __init__(self, device="cpu", **kwargs):
        super().__init__()
        resnet = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
        self.resnet = nn.Sequential(*[i for i in list(resnet.children())[:-1]]).to(device)
        for ch in self.resnet.children():
            for param in ch.parameters():
                param.requires_grad = False
        self.gat = GAT(**kwargs).to(device)
        self.device = device
        self.adj = torch.tensor(grid_to_graph(7, 7, return_as=np.ndarray)).to(device)
        self.avg_pool = nn.AvgPool2d(kernel_size=8, stride=8, padding=0, ceil_mode=False,
                                               count_include_pad=False)
        
    def forward(self, x):
        x = self.resnet(x.to(self.device))
        x = self.avg_pool(x)
        #x = x.view(512, -1)
        #x = x.view(512, 1)
        x = self.gat(x, self.adj)
        return torch.sigmoid(x)
        


preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

torch.manual_seed(42)
kwargs = {"nfeat":512, "nhid":64, "nclass":1, "nheads":49, "dropout":0.6, "alpha":0.01}
model = AntispoofModel(**kwargs, device="cpu")
test_dataset = ASDataset(client_file="raw/client_train_raw.txt", imposter_file="raw/imposter_train_raw.txt", \
    transforms=preprocess)
train_dataset = ASDataset(client_file="raw/client_test_raw.txt", imposter_file="raw/imposter_test_raw.txt", \
    transforms=preprocess)
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=True)
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.3)

def forward(self, x):
     48         x = self.resnet(x.to(self.device))
---> 49         x = self.avg_pool(x)
     50         #x = x.view(512, -1)
     51         #x = x.view(512, 1)

~\anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1188         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1189                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1190             return forward_call(*input, **kwargs)
   1191         # Do not call functions when jit is used
   1192         full_backward_hooks, non_full_backward_hooks = [], []

~\anaconda3\lib\site-packages\torch\nn\modules\pooling.py in forward(self, input)
    625 
    626     def forward(self, input: Tensor) -> Tensor:
--> 627         return F.avg_pool2d(input, self.kernel_size, self.stride,
    628                             self.padding, self.ceil_mode, self.count_include_pad, self.divisor_override)
    629 

RuntimeError: Given input size: (512x1x1). Calculated output size: (512x0x0). Output size is too small

При удалении avgpool

class AntispoofModel(nn.Module):
    def __init__(self, device="cpu", **kwargs):
        super().__init__()
        resnet = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
        self.resnet = nn.Sequential(*[i for i in list(resnet.children())[:-1]]).to(device)
        for ch in self.resnet.children():
            for param in ch.parameters():
                param.requires_grad = False
        self.gat = GAT(**kwargs).to(device)
        self.device = device
        self.adj = torch.tensor(grid_to_graph(7, 7, return_as=np.ndarray)).to(device)
        self.avg_pool = nn.AvgPool2d(kernel_size=1, stride=1, padding=0, ceil_mode=False,
                                               count_include_pad=False)
        
    def forward(self, x):
        x = self.resnet(x.to(self.device))
        #x = self.avg_pool(x)
        x = x.view(512, -1)
        #x = x.view(512, 1)
        x = self.gat(x, self.adj)
        return torch.sigmoid(x)
        


preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

torch.manual_seed(42)
kwargs = {"nfeat":512, "nhid":64, "nclass":1, "nheads":49, "dropout":0.6, "alpha":0.01}
model = AntispoofModel(**kwargs, device="cpu")

train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=True)
criterion = nn.BCELoss()

~\AppData\Local\Temp\ipykernel_2940\480930026.py in forward(self, x)
     50         x = x.view(512, -1)
     51         #x = x.view(512, 1)
---> 52         x = self.gat(x, self.adj)
     53         return torch.sigmoid(x)
     54 

~\anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1188         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1189                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1190             return forward_call(*input, **kwargs)
   1191         # Do not call functions when jit is used
   1192         full_backward_hooks, non_full_backward_hooks = [], []

~\Desktop\antispoof\GAT\models.py in forward(self, x, adj)
     21     def forward(self, x, adj):
     22         x = F.dropout(x, self.dropout, training=self.training)
---> 23         x = torch.cat([att(x, adj, i) for i, att in enumerate(self.attentions)], dim=2)
     24         x = F.dropout(x, self.dropout, training=self.training)
     25         x = F.elu(self.out_att(x, adj))

~\Desktop\antispoof\GAT\models.py in <listcomp>(.0)
     21     def forward(self, x, adj):
     22         x = F.dropout(x, self.dropout, training=self.training)
---> 23         x = torch.cat([att(x, adj, i) for i, att in enumerate(self.attentions)], dim=2)
     24         x = F.dropout(x, self.dropout, training=self.training)
     25         x = F.elu(self.out_att(x, adj))

~\anaconda3\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
   1188         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1189                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1190             return forward_call(*input, **kwargs)
   1191         # Do not call functions when jit is used
   1192         full_backward_hooks, non_full_backward_hooks = [], []

~\Desktop\antispoof\GAT\layers.py in forward(self, h, adj, i)
     25 
     26     def forward(self, h, adj, i=None):
---> 27         Wh = torch.matmul(h, self.W) # h.shape: (N, in_features), Wh.shape: (N, out_features)
     28         e = self._prepare_attentional_mechanism_input(Wh)
     29         zero_vec = -9e15*torch.ones_like(e)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (512x64 and 512x64)