Bugfix: Initialize grad_input, grad_weight to None (#68)

* initialize grad_input to None

* minor

axonn/intra_layer/fully_connected.py

@@ -107,6 +107,8 @@ def backward(ctx, grad_output):
         if dist.get_world_size(ctx.backward_all_reduce_group) > 1 or (
             not overlap_reduce_scatter
         ):
+            grad_input, grad_weight = None, None
+
             if ctx.needs_input_grad[0]:
                 grad_input = grad_output.matmul(weight)
                 handle = dist.all_reduce(
@@ -136,6 +138,8 @@ def backward(ctx, grad_output):
                 grad_weight = None  # weight gradients are not ready yet
             return grad_input, grad_weight, None, None, None, None, None, None, None
         else:
+            grad_input, grad_weight = None, None
+
             if ctx.needs_input_grad[1]:
                 grad_weight = (
                     grad_output.reshape(-1, grad_output.shape[-1])