Merge pull request #686 from dcslin/linear-layer-input-channel-check

Linear layer should check if the input channels parameter match inputs
apache · Apr 28, 2020 · 12efad4 · 12efad4
2 parents 893992c + ac0b8e8
commit 12efad4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/singa/autograd.py b/python/singa/autograd.py
@@ -1409,6 +1409,9 @@ def __call__(self, x):
             self.device_check(x, self.W, self.b)
         else:
             self.device_check(x, self.W)
+        assert x.shape[1] == self.W.shape[0], (
+            "Linear layer expects input features size %d received %d" %
+            (self.W.shape[0], x.shape[1]))
         y = matmul(x, self.W)
         if self.bias:
             y = add_bias(y, self.b, axis=0)