| 45 | |
| 46 | class AdaptiveBatchNorm2d(nn.Module): |
| 47 | def __init__(self, num_features, eps=1e-5, momentum=0.1): |
| 48 | super().__init__() |
| 49 | self.num_features = num_features |
| 50 | self.eps = eps |
| 51 | self.momentum = momentum |
| 52 | |
| 53 | self.weight = nn.Parameter(torch.ones(num_features)) |
| 54 | self.bias = nn.Parameter(torch.zeros(num_features)) |
| 55 | self.register_buffer('running_mean', torch.zeros(num_features)) |
| 56 | self.register_buffer('running_var', torch.ones(num_features)) |
| 57 | self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long)) |
| 58 | |
| 59 | def forward(self, x): |
| 60 | if self.training: |