Skip to content

GitLab

Unverified Commit 07e42152 authored by Jerry Jiarui XU's avatar Jerry Jiarui XU Committed by GitHub
Browse files

Fixed group of syncbn (#427)

parent 4a044c64
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
mmcv/ops/sync_bn.py
View file @ 07e42152
...@@ -126,7 +126,8 @@ class SyncBatchNorm(Module): ...@@ -126,7 +126,8 @@ class SyncBatchNorm(Module):
self.momentum = momentum self.momentum = momentum
self.affine = affine self.affine = affine
self.track_running_stats = track_running_stats self.track_running_stats = track_running_stats
self.group = dist.group.WORLD if group is None else group group = dist.group.WORLD if group is None else group
self.group = group
self.group_size = dist.get_world_size(group) self.group_size = dist.get_world_size(group)
if self.affine: if self.affine:
self.weight = Parameter(torch.Tensor(num_features)) self.weight = Parameter(torch.Tensor(num_features))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment