Accommodate johnmarktaylor91/torchlens#18

Author: kalekundert

escnn/nn/modules/pooling/gaussian_blur.py

@@ -10,11 +10,6 @@
 
 from typing import Optional, Union, Tuple
 
-_CONV = {
-        2: F.conv2d,
-        3: F.conv3d,
-}
-
 class GaussianBlurND(LazyModuleMixin, Module):
 
     def __init__(
@@ -76,6 +71,7 @@ def __init__(
         self.stride = stride
         self.edge_correction = edge_correction
         self.d = d
+        self.conv = getattr(F, f'conv{d}d')
 
         if padding is not None:
             self.padding = padding
@@ -130,7 +126,7 @@ def forward(self, x):
         return y
 
     def blur(self, x):
-        return _CONV[self.d](
+        return self.conv(
                 x,
                 self.filter,
                 stride=self.stride,