Fixed bug in block sparse linear backward function.

Author: madlag

pytorch_block_sparse/block_sparse_linear.py

@@ -110,8 +110,11 @@ def backward(ctx, grad_output):
         else:
             grad_weight1 = None
 
-        assert(not (grad_weight1 == 0).all())
-        assert(grad_input1.shape == input.shape)
+        if grad_weight1 != None:
+            assert(not (grad_weight1 == 0).all())
+        if grad_input1 != None:
+            assert(grad_input1.shape == input.shape)
+
         return grad_input1, grad_weight1, None
 
 class BlockSparseLinear(nn.Module):