You could write a custom nn.Module and apply the squeeze operation there.
Something like this should work:
class Squeeze(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        
    def forward(self, x):
        x = x.squeeze(dim=self.dim)
        return x

squeeze = Squeeze(dim=2)
x = torch.randn(18, 18, 1, 512)
out = squeeze(x)
print(out.shape)
> torch.Size([18, 18, 512])