pytorch-Flatten操作

class Flatten(nn.Module):
    def __init__(self):
        super(Flatten,self).__init__()
        
    def forward(self,input):
        shape = torch.prod(torch.tensor(x.shape[1:])).item()
        # -1 把第一个维度保持住
        return x.view(-1,shape)