Merge pull request #974 from lscheinkman/fix_batchnorm

Fix sparsenet with batchnorm
numenta · Mar 22, 2019 · 77b3c0f · 77b3c0f
2 parents fc5cb7e + 136b5a2
commit 77b3c0f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/htmresearch/frameworks/pytorch/sparse_net.py b/htmresearch/frameworks/pytorch/sparse_net.py
@@ -245,7 +245,7 @@ def __init__(self,
 
         # Batch Norm
         if useBatchNorm:
-          bn = nn.BatchNorm2d(outChannels, affine=False)
+          bn = nn.BatchNorm2d(outChannels[i], affine=False)
           cnnSdr.add_module("cnnSdr{}_bn".format(i + 1), bn)
 
         # Max pool