Speed profiling improvements (#2648)

* Speed profiling improvements

* Update torch_utils.py

deepcopy() required to avoid adding elements to model.

* Update torch_utils.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-03-29 15:19:07 +02:00
1 parent 1e8ab3f5f2
commit 866bc7d640
2 files changed
+5 -4

No files matched your search

+1 -1
View File
@@ -191,7 +191,7 @@ def fuse_conv_and_bn(conv, bn):
# prepare filters
w_conv = conv.weight.clone().view(conv.out_channels, -1)
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
# prepare spatial bias
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias