Add torchvision.ops.FrozenBatchNorm2d to types

chr5tphr · Oct 4, 2023 · fc915a0 · fc915a0
1 parent 60a2c08
commit fc915a0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/zennit/types.py b/src/zennit/types.py
@@ -17,7 +17,7 @@
 # along with this library. If not, see <https://www.gnu.org/licenses/>.
 '''Type definitions for convenience.'''
 import torch
-
+import torchvision
 
 class SubclassMeta(type):
     '''Meta class to bundle multiple subclasses.'''
@@ -71,6 +71,7 @@ class BatchNorm(metaclass=SubclassMeta):
         torch.nn.modules.batchnorm.BatchNorm1d,
         torch.nn.modules.batchnorm.BatchNorm2d,
         torch.nn.modules.batchnorm.BatchNorm3d,
+        torchvision.ops.FrozenBatchNorm2d,
     )