From 4bb7eb8b849fc8a90823a60e2b7a8ec9e38926bf Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 12 Apr 2022 11:02:11 +0200
Subject: [PATCH] Dynamic normalization layer selection (#7392)

* Dynamic normalization layer selection

Based on actual available layers. Torch 1.7 compatible, resolves https://github.com/ultralytics/yolov5/issues/7381

* Update train.py
---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 80bff18fd653..806e2cebe561 100644
--- a/train.py
+++ b/train.py
@@ -151,7 +151,7 @@ def train(hyp, opt, device, callbacks):  # hyp is path/to/hyp.yaml or hyp dictio
     LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
 
     g = [], [], []  # optimizer parameter groups
-    bn = nn.BatchNorm2d, nn.LazyBatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d, nn.LazyInstanceNorm2d, nn.LayerNorm
+    bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k)  # normalization layers, i.e. BatchNorm2d()
     for v in model.modules():
         if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias
             g[2].append(v.bias)