feat: enable BatchNorm2d in CNNEncoder and add AMP support

- Uncomment BatchNorm2d in CNNEncoder (activated when cfg.use_bn=True) - Add torch.amp.GradScaler + autocast for mixed precision training - Add --amp/--no-amp CLI argument (default: enabled) Generated by Mistral Vibe. deepseek-v4-flash Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
2026-06-08 00:41:34 +08:00
parent b5abbc239d
commit 1369edaad7
2 changed files with 23 additions and 12 deletions
@@ -29,7 +29,7 @@ class CNNEncoder(nn.Module):
        for out_ch in channels:
            layers.extend([
                nn.Conv2d(in_ch, out_ch, kernel_size=cfg.kernel_size, padding=cfg.kernel_size // 2),
-                # nn.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(),
+                nn.BatchNorm2d(out_ch) if cfg.use_bn else nn.Identity(),
                nn.Identity(),
                nn.LeakyReLU(inplace=True),
                nn.MaxPool2d(cfg.pool_size),