Accommodate johnmarktaylor91/torchlens#18

Signed-off-by: Kale Kundert <[email protected]>
kalekundert · Nov 21, 2023 · 0d230db · 0d230db
1 parent 948307d
commit 0d230db
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/escnn/nn/modules/pooling/gaussian_blur.py b/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,