I’m a bit confused about the error message and the information in your post.
While the error message suggests you are using nn.AvgPool3d and are passing a wrong input shape, you are mentioning nn.AvgPool2d for the posted input.
Anyway, these codes should work:
# 3D
x = torch.randn([1, 2, 3, 256, 256])
pool = nn.AvgPool3d(kernel_size=3, stride=2, padding=1)
out = pool(x)

# 2D
x = torch.randn([2, 3, 256, 256])
pool = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
out = pool(x)