Adding more norm layers in split_normalization_params.

Author: datumbox

torchvision/ops/_utils.py

@@ -3,6 +3,8 @@
 import torch
 from torch import nn, Tensor
 
+from .misc import FrozenBatchNorm2d
+
 
 def _cat(tensors: List[Tensor], dim: int = 0) -> Tensor:
     """
@@ -43,7 +45,13 @@ def split_normalization_params(
 ) -> Tuple[List[Tensor], List[Tensor]]:
     # Adapted from https://github.com/facebookresearch/ClassyVision/blob/659d7f78/classy_vision/generic/util.py#L501
     if not norm_classes:
-        norm_classes = [nn.modules.batchnorm._BatchNorm, nn.LayerNorm, nn.GroupNorm]
+        norm_classes = [
+            nn.modules.batchnorm._BatchNorm,
+            nn.LayerNorm,
+            nn.GroupNorm,
+            nn.modules.instancenorm._InstanceNorm,
+            nn.LocalResponseNorm,
+        ]
 
     for t in norm_classes:
         if not issubclass(t, nn.Module):