corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Batch norm tracking

Browse source
This commit is contained in:
Corentin Risselin 2020-07-07 11:04:57 +09:00
commit 7f4a162033
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
layers.py
View file

@ -68,7 +68,7 @@ class Linear(Layer):
self.batch_norm = nn.BatchNorm1d( self.batch_norm = nn.BatchNorm1d(
out_channels, out_channels,
momentum=Layer.BATCH_NORM_MOMENTUM, momentum=Layer.BATCH_NORM_MOMENTUM,
track_running_stats=not Layer.BATCH_NORM_TRAINING) if self.batch_norm else None track_running_stats=Layer.BATCH_NORM_TRAINING) if self.batch_norm else None
def forward(self, input_data: torch.Tensor) -> torch.Tensor: def forward(self, input_data: torch.Tensor) -> torch.Tensor:
return super().forward(self.fc(input_data)) return super().forward(self.fc(input_data))